added prune_dict function

basic_utils/helpers.py

@@ -165,3 +165,15 @@ def dedupe(items, key=None):
         if val not in seen:
             yield item
             seen.add(val)
+
+
+def prune_dict(d, key=lambda x: x is not None):
+    """
+    Returns new dictionary with key / values pairs filtered by key function.
+    Prunes None values by default
+
+    >>> d = {'Homer': 39, 'Marge': 36, 'Bart': 10}
+    >>> prune_dict(d, key=lambda x: x < 20)
+    {'Bart': 10}
+    """
+    return {k: v for k, v in d.items() if key(v)}

tests/test_helpers.py

@@ -1,7 +1,6 @@
 import pytest
 
 from basic_utils.helpers import (
-    uniq,
     butlast,
     dedupe,
     dict_subset,
@@ -11,9 +10,11 @@
     get_keys,
     last,
     partial_flatten,
+    prune_dict,
     rest,
     reverse,
     set_in_dict,
+    uniq,
 )
 
 
@@ -148,3 +149,11 @@ def test_uniq(data, expected):
 def test_dedupe(data, key, expected):
     """Tests that dedupe removes duplicates whilst preserving order"""
     assert list(dedupe(data, key)) == expected
+
+
+@pytest.mark.parametrize("args, expected", [
+    (({'Homer': 39, 'Marge': 36, 'Bart': 10}, lambda x: x < 20), {'Bart': 10}),
+    (({'a': None, 'b': 2, 'c': None},), {'b': 2})
+])
+def test_prune_dict(args, expected):
+    assert prune_dict(*args) == expected