Add ignore_key_func

flatten_json/__init__.py

@@ -33,7 +33,10 @@ def _construct_key(previous_key, separator, new_key):
         return new_key
 
 
-def flatten(nested_dict, separator="_", root_keys_to_ignore=set()):
+def flatten(nested_dict,
+            separator="_",
+            root_keys_to_ignore=set(),
+            ignore_key_func=None):
     """
     Flattens a dictionary with nested structure to a dictionary with no
     hierarchy
@@ -44,6 +47,7 @@ def flatten(nested_dict, separator="_", root_keys_to_ignore=set()):
     :param nested_dict: dictionary we want to flatten
     :param separator: string to separate dictionary keys by
     :param root_keys_to_ignore: set of root keys to ignore from flattening
+    :param ignore_key_func: func to ignore key at any level from flattening
     :return: flattened dictionary
     """
     assert isinstance(nested_dict, dict), "flatten requires a dictionary input"
@@ -68,10 +72,13 @@ def _flatten(object_, key):
         # These object types support iteration
         elif isinstance(object_, dict):
             for object_key in object_:
-                if not (not key and object_key in root_keys_to_ignore):
-                    _flatten(object_[object_key], _construct_key(key,
-                                                                 separator,
-                                                                 object_key))
+                if not key and object_key in root_keys_to_ignore:
+                    continue
+                if ignore_key_func and ignore_key_func(object_key):
+                    continue
+                _flatten(object_[object_key], _construct_key(key,
+                                                             separator,
+                                                             object_key))
         elif isinstance(object_, (list, set, tuple)):
             for index, item in enumerate(object_):
                 _flatten(item, _construct_key(key, separator, index))

test_flatten.py

@@ -2175,14 +2175,15 @@ def test_unflatten_with_list_deep(self):
     def test_flatten_ignore_keys(self):
         """Ignore a set of root keys for processing"""
         dic = {
-            'a': {'a': [1, 2, 3]},
+            'a': {'a': [1, 2, 3], 'b': 4},
             'b': {'b': 'foo', 'c': 'bar'},
             'c': {'c': [{'foo': 5, 'bar': 6, 'baz': [1, 2, 3]}]}
         }
         expected = {
             'a_a_0': 1,
             'a_a_1': 2,
-            'a_a_2': 3
+            'a_a_2': 3,
+            'a_b': 4
         }
         actual = flatten(dic, root_keys_to_ignore={'b', 'c'})
         self.assertEqual(actual, expected)
@@ -2195,6 +2196,23 @@ def test_command_line(self):
         result = json.loads(output)
         self.assertEqual(result, dict(a_b=1))
 
+    def test_keys_to_ignore(self):
+        """Ignore a set of keys any depth for processing"""
+        dic = {
+            'a': {'a': [1, 2, 3], 'b': 4},
+            'b': {'b': 'foo', 'c': 'bar'}
+        }
+        expected = {
+            'a_a_0': 1,
+            'a_a_1': 2,
+            'a_a_2': 3
+        }
+
+        def _ignore_key_func(key):
+            return key == 'b'
+        actual = flatten(dic, ignore_key_func=_ignore_key_func)
+        self.assertEqual(actual, expected)
+
 
 if __name__ == '__main__':
     unittest.main()