Implement omit and pick dicts functions. (#304)

lib/dicts.bzl

@@ -38,6 +38,33 @@ def _add(*dictionaries, **kwargs):
     result.update(kwargs)
     return result
 
+def _omit(dictionary, keys):
+    """Returns a new `dict` that has all the entries of `dictionary` with keys not in `keys`.
+
+    Args:
+      dictionary: A `dict`.
+      keys: A sequence.
+
+    Returns:
+      A new `dict` that has all the entries of `dictionary` with keys not in `keys`.
+    """
+    keys_set = {k: None for k in keys}
+    return {k: dictionary[k] for k in dictionary if k not in keys_set}
+
+def _pick(dictionary, keys):
+    """Returns a new `dict` that has all the entries of `dictionary` with keys in `keys`.
+
+    Args:
+      dictionary: A `dict`.
+      keys: A sequence.
+
+    Returns:
+      A new `dict` that has all the entries of `dictionary` with keys in `keys`.
+    """
+    return {k: dictionary[k] for k in keys if k in dictionary}
+
 dicts = struct(
     add = _add,
+    omit = _omit,
+    pick = _pick,
 )

tests/dicts_tests.bzl

@@ -83,9 +83,77 @@ def _add_test(ctx):
 
 add_test = unittest.make(_add_test)
 
+def _omit_test(ctx):
+    """Unit tests for dicts.omit."""
+    env = unittest.begin(ctx)
+
+    # Test empty dict, empty list.
+    asserts.equals(env, {}, dicts.omit({}, []))
+
+    # Test empty dict, nonempty list.
+    asserts.equals(env, {}, dicts.omit({}, ["a"]))
+
+    # Test nonempty dict, empty list.
+    asserts.equals(env, {"a": 1}, dicts.omit({"a": 1}, []))
+
+    # Test key in dict.
+    asserts.equals(env, {}, dicts.omit({"a": 1}, ["a"]))
+
+    # Test key not in dict.
+    asserts.equals(env, {"a": 1}, dicts.omit({"a": 1}, ["b"]))
+
+    # Since dictionaries are passed around by reference, make sure that the
+    # result of dicts.omit is always a *copy* by modifying it afterwards and
+    # ensuring that the original argument doesn't also reflect the change. We do
+    # this to protect against someone who might attempt to optimize the function
+    # by returning the argument itself in the empty list case.
+    original = {"a": 1}
+    result = dicts.omit(original, [])
+    result["a"] = 2
+    asserts.equals(env, 1, original["a"])
+
+    return unittest.end(env)
+
+omit_test = unittest.make(_omit_test)
+
+def _pick_test(ctx):
+    """Unit tests for dicts.pick."""
+    env = unittest.begin(ctx)
+
+    # Test empty dict, empty list.
+    asserts.equals(env, {}, dicts.pick({}, []))
+
+    # Test empty dict, nonempty list.
+    asserts.equals(env, {}, dicts.pick({}, ["a"]))
+
+    # Test nonempty dict, empty list.
+    asserts.equals(env, {}, dicts.pick({"a": 1}, []))
+
+    # Test key in dict.
+    asserts.equals(env, {"a": 1}, dicts.pick({"a": 1}, ["a"]))
+
+    # Test key not in dict.
+    asserts.equals(env, {}, dicts.pick({"a": 1}, ["b"]))
+
+    # Since dictionaries are passed around by reference, make sure that the
+    # result of dicts.pick is always a *copy* by modifying it afterwards and
+    # ensuring that the original argument doesn't also reflect the change. We do
+    # this to protect against someone who might attempt to optimize the function
+    # by returning the argument itself.
+    original = {"a": 1}
+    result = dicts.pick(original, ["a"])
+    result["a"] = 2
+    asserts.equals(env, 1, original["a"])
+
+    return unittest.end(env)
+
+pick_test = unittest.make(_pick_test)
+
 def dicts_test_suite():
     """Creates the test targets and test suite for dicts.bzl tests."""
     unittest.suite(
         "dicts_tests",
         add_test,
+        omit_test,
+        pick_test,
     )