Merge pull request #1085 from TanYewWei/issue_1069

Issue #1069: Add Dict.drop, Dict.take and Dict.split

Author: José Valim

lib/elixir/lib/dict.ex

@@ -41,9 +41,11 @@ defmodule Dict do
 
   @type key :: any
   @type value :: any
+  @type keys :: [ key ]
   @type t :: tuple | list
 
   defcallback delete(t, key) :: t
+  defcallback drop(t, keys) :: t
   defcallback empty(t) :: t
   defcallback equal?(t, t) :: boolean
   defcallback get(t, key) :: value
@@ -58,6 +60,8 @@ defmodule Dict do
   defcallback put(t, key, value) :: t
   defcallback put_new(t, key, value) :: t
   defcallback size(t) :: non_neg_integer()
+  defcallback split(t, keys) :: {t, t}
+  defcallback take(t, keys) :: t
   defcallback to_list(t) :: list()
   defcallback update(t, key, (value -> value)) :: t | no_return
   defcallback update(t, key, value, (value -> value)) :: t
@@ -362,6 +366,83 @@ defmodule Dict do
     target(dict).update(dict, key, initial, fun)
   end
 
+  @doc """
+  
+  Returns a tuple of two dicts, where the first dict contains only 
+  entries from `dict` with keys in `keys`, and the second dict 
+  contains only entries from `dict` with keys not in `keys`
+ 
+  Any non-member keys are ignored.
+
+  ## Examples 
+
+      iex> d = HashDict.new([a: 1, b: 2])
+      ...> { d1, d2 } = Dict.split(d, [:a, :c])
+      ...> { Dict.to_list(d1), Dict.to_list(d2) }
+      { [a: 1], [b: 2] }
+
+      iex> d = HashDict.new([])
+      ...> { d1, d2 } = Dict.split(d, [:a, :c])
+      ...> { Dict.to_list(d1), Dict.to_list(d2) }
+      { [], [] }
+
+      iex> d = HashDict.new([a: 1, b: 2])
+      ...> { d1, d2 } = Dict.split(d, [:a, :b, :c])
+      ...> { Dict.to_list(d1), Dict.to_list(d2) }
+      { [a: 1, b: 2], [] }
+
+  """
+  @spec split(t, keys) :: {t, t}
+  def split(dict, keys) do
+    target(dict).split(dict, keys)
+  end
+
+  @doc """
+  
+  Returns a new dict where the the given `keys` a removed from `dict`.
+  Any non-member keys are ignored.
+
+  ## Examples 
+
+      iex> d = HashDict.new([a: 1, b: 2])
+      ...> d = Dict.drop(d, [:a, :c, :d])
+      ...> Dict.to_list(d)
+      [b: 2]
+
+      iex> d = HashDict.new([a: 1, b: 2])
+      ...> d = Dict.drop(d, [:c, :d])
+      ...> Dict.to_list(d)
+      [a: 1, b: 2]
+
+  """
+  @spec drop(t, keys) :: t
+  def drop(dict, keys) do
+    target(dict).drop(dict, keys)
+  end
+
+  @doc """
+  
+  Returns a new dict where only the keys in `keys` from `dict` are
+  included. Any non-member keys are ignored.
+
+  ## Examples 
+
+      iex> d = HashDict.new([a: 1, b: 2])
+      ...> d = Dict.take(d, [:a, :c, :d])
+      ...> Dict.to_list(d)
+      [a: 1]
+
+      iex> d = HashDict.new([a: 1, b: 2])
+      ...> d = Dict.take(d, [:c, :d])
+      ...> Dict.to_list(d)
+      []
+
+  """
+  @spec take(t, keys) :: t
+  def take(dict, keys) do
+    target(dict).take(dict, keys)
+  end  
+
   @doc """
   Returns an empty dict of the same type as `dict`.
   """

lib/elixir/lib/hash_dict.ex

@@ -286,6 +286,53 @@ defmodule HashDict do
     end
   end
 
+  @doc """
+  Splits a dict into two dicts, 
+  one containing entries with key in the keys list, 
+  and another containing entries with key not in keys.
+  Returns a 2-tuple of the new dicts.
+  """
+  def split(dict, keys) when keys == [] do
+    { new(), dict }
+  end
+
+  def split(dict, keys) do
+    acc = { new(), new() }
+    dict_fold dict, acc, fn({ k, v }, { take, drop }) ->
+      if Enum.member?(keys, k) do
+        { put(take, k, v), drop }
+      else 
+        { take, put(drop, k, v) }
+      end
+    end
+  end
+
+  @doc """
+  Returns a new dict with only the entries
+  which key is in keys
+  """
+  def take(_, keys) when keys == [] do
+    new()
+  end
+
+  def take(dict, keys) do
+    { members, _ } = split(dict, keys)
+    members
+  end
+
+  @doc """
+  Returns a new dict with only the entries 
+  which key is not in keys
+  """
+  def drop(dict, keys) when keys == [] do 
+    dict
+  end
+
+  def drop(dict, keys) do
+    { _, non_members } = split(dict, keys)
+    non_members
+  end
+
   ## Dict-wide functions
 
   defp dict_get(ordered(bucket: bucket), key) do

lib/elixir/lib/list_dict.ex

@@ -140,6 +140,53 @@ defmodule ListDict do
     end
   end
 
+  @doc """
+  Splits a dict into two dicts, 
+  one containing entries with key in the keys list, 
+  and another containing entries with key not in keys.
+  Returns a 2-tuple of the new dicts.
+  """
+  def split(dict, keys) when keys == [] do
+    { new(), dict }
+  end
+
+  def split(dict, keys) do
+    acc = { new(), new() }
+    Enum.reduce dict, acc, fn({ k, v }, { take, drop }) ->
+      if Enum.member?(keys, k) do
+        { [{k,v}|take], drop }
+      else 
+        { take, [{k,v}|drop] }
+      end
+    end
+  end
+
+  @doc """
+  Returns a new dict with only the entries
+  which key is in keys
+  """
+  def take(_, keys) when keys == [] do
+    new()
+  end
+
+  def take(dict, keys) do
+    { members, _ } = split(dict, keys)
+    members
+  end
+
+  @doc """
+  Returns a new dict with only the entries 
+  which key is not in keys
+  """
+  def drop(dict, keys) when keys == [] do 
+    dict
+  end
+
+  def drop(dict, keys) do
+    { _, non_members } = split(dict, keys)
+    non_members
+  end
+
   @doc """
   Updates the key in the dict according to the given function.
   """

lib/elixir/test/elixir/dict_test.exs

@@ -160,6 +160,49 @@ defmodule DictTest.Common do
         assert dicts_equal actual, new_dict
       end
 
+			test :split do
+				split_keys = []
+				{take, drop} = Dict.split(new_dict, split_keys)
+				assert dicts_equal take, Dict.empty(new_dict)
+				assert dicts_equal drop, new_dict
+				
+				split_keys = ["unknown_key"]
+				{take, drop} = Dict.split(new_dict, split_keys)
+				assert dicts_equal take, Dict.empty(new_dict)
+				assert dicts_equal drop, new_dict
+
+				split_keys = ["first_key", "second_key", "unknown_key"]
+				{take, drop} = Dict.split(new_dict, split_keys)				
+				first_val = Dict.get(new_dict, "first_key")
+				second_val = Dict.get(new_dict, "second_key")
+				take_expected = Dict.empty(new_dict)
+				take_expected = Dict.put(take_expected, "first_key", first_val)
+				take_expected = Dict.put(take_expected, "second_key", second_val)			
+				drop_expected = new_dict
+				drop_expected = Dict.delete(drop_expected, "first_key")
+				drop_expected = Dict.delete(drop_expected, "second_key")
+				assert dicts_equal take, take_expected
+				assert dicts_equal drop, drop_expected
+			end
+
+			test :take do
+				result = Dict.take(new_dict, ["unknown_key"])
+				assert dicts_equal result, Dict.empty(new_dict)
+
+				result = Dict.take(new_dict, ["first_key"])
+				first_val = Dict.get(new_dict, "first_key")
+				expected = Dict.put(Dict.empty(new_dict), "first_key", first_val)
+				assert dicts_equal result, expected
+			end
+
+			test :drop do
+				result = Dict.drop(new_dict, ["unknown_key"])
+				assert dicts_equal result, new_dict
+
+				result = Dict.drop(new_dict, ["first_key"])
+				assert dicts_equal result, Dict.delete(new_dict, "first_key")
+			end
+
       test :empty do
         assert empty_dict == Dict.empty new_dict
       end