diff --git a/flatten_json/__init__.py b/flatten_json/__init__.py index a1497c9..2935b9a 100644 --- a/flatten_json/__init__.py +++ b/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)) diff --git a/test_flatten.py b/test_flatten.py index 10b38cc..885e342 100644 --- a/test_flatten.py +++ b/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()