diff --git a/doc/source/api.rst b/doc/source/api.rst index 7c1ff662..e6e887f6 100644 --- a/doc/source/api.rst +++ b/doc/source/api.rst @@ -84,6 +84,7 @@ Dicttoolz keymap merge merge_with + select_keys update_in valfilter valmap diff --git a/toolz/dicttoolz.py b/toolz/dicttoolz.py index 457bc269..94c268ec 100644 --- a/toolz/dicttoolz.py +++ b/toolz/dicttoolz.py @@ -337,3 +337,22 @@ def get_in(keys, coll, default=None, no_default=False): if no_default: raise return default + + +def select_keys(keys, d, factory=dict): + """ Select only certain keys from a dictionary + + If supplied keys are not found in the dictionary, returns an empty + dictionary. + + >>> d = {'a': 1, 'b': 2, 'c': 3} + >>> select_keys(['a', 'b', 'd'], d) # doctest: +SKIP + {'a': 1, 'b': 2} + >>> select_keys(['d'], d) + {} + """ + rv = factory() + for k in keys: + if k in d: + rv[k] = d[k] + return rv diff --git a/toolz/tests/test_dicttoolz.py b/toolz/tests/test_dicttoolz.py index d45cd6cf..ff060d94 100644 --- a/toolz/tests/test_dicttoolz.py +++ b/toolz/tests/test_dicttoolz.py @@ -1,10 +1,11 @@ +import os from collections import defaultdict as _defaultdict from collections.abc import Mapping -import os -from toolz.dicttoolz import (merge, merge_with, valmap, keymap, update_in, - assoc, dissoc, keyfilter, valfilter, itemmap, - itemfilter, assoc_in) from toolz.functoolz import identity +from toolz.dicttoolz import (assoc, assoc_in, dissoc, itemfilter, + itemmap, keyfilter, keymap, merge, + merge_with, select_keys, update_in, + valfilter, valmap) from toolz.utils import raises @@ -119,6 +120,14 @@ def test_assoc_in(self): assert d is oldd assert d2 is not oldd + def test_select_keys(self): + D, kw = self.D, self.kw + assert select_keys([], D({}), **kw) == D({}) + assert select_keys(["a"], D({"a": 1}), **kw) == D({"a": 1}) + assert select_keys([], D({"a": 1}), **kw) == D({}) + assert select_keys(["b"], D({"a": 1}), **kw) == D({}) + assert select_keys(["b"], D({"a": 1, "b": 2}), **kw) == D({"b": 2}) + def test_update_in(self): D, kw = self.D, self.kw assert update_in(D({"a": 0}), ["a"], inc, **kw) == D({"a": 1})