diff --git a/lib/dicts.bzl b/lib/dicts.bzl index 3f8e661..5a92aa1 100644 --- a/lib/dicts.bzl +++ b/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, ) diff --git a/tests/dicts_tests.bzl b/tests/dicts_tests.bzl index 2fc407d..816fcea 100644 --- a/tests/dicts_tests.bzl +++ b/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, )