From 68150e4f74d902b00c80c0da23a46acdc4a435ad Mon Sep 17 00:00:00 2001 From: Bernardo Gomes Date: Fri, 1 Jul 2022 15:52:47 -0300 Subject: [PATCH] Add suport to list --- jsonmask/mask.py | 17 +++++++++++++++-- jsonmask/tests/test_mask.py | 3 +++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/jsonmask/mask.py b/jsonmask/mask.py index ce09c0e..155736a 100644 --- a/jsonmask/mask.py +++ b/jsonmask/mask.py @@ -36,14 +36,27 @@ def apply_json_mask(data, json_mask, is_negated=False, depth=1, max_depth=None): for key, subdata in data.items(): if should_include_variable(key, json_mask, is_negated=is_negated): + next_json_mask = json_mask.get(key, {}) + if isinstance(subdata, list): + allowed_data[key] = [ + apply_json_mask( + entry, + next_json_mask, + is_negated=is_negated, + depth=depth + 1, + max_depth=max_depth, + ) + if isinstance(entry, dict) + else entry + for entry in subdata + ] + continue # Terminal data if not isinstance(subdata, dict): allowed_data[key] = subdata continue - next_json_mask = json_mask.get(key, {}) - # Dead ends in the mask indicate that want # everything nested below this if not next_json_mask: diff --git a/jsonmask/tests/test_mask.py b/jsonmask/tests/test_mask.py index aea999d..13b9833 100644 --- a/jsonmask/tests/test_mask.py +++ b/jsonmask/tests/test_mask.py @@ -23,6 +23,9 @@ def test_apply_json_mask(): ({'a': {'b': 1}, 'b': {'asdf': 2}}, 'a,b/c', {'a': {'b': 1}, 'b': {}}, {'b': {'asdf': 2}},), ({'a': {'b': 1}, 'b': {'asdf': 2}}, 'a,b/asdf', ORIGINAL, {'b': {}},), ({'a': {'b': 1}, 'b': {'asdf': 2}}, 'a,c/c', {'a': {'b': 1}}, {'b': {'asdf': 2}},), + ({'a': [1,2,3]}, 'a', ORIGINAL, EMPTY,), + ({'a': [{'b': 1, 'c': [1,2,3]}]}, 'a(b)', {'a': [{'b':1}]}, {'a': [{'c': [1,2,3]}]},), + ({'a': [{'b': 1, 'c': [{'d': 1}, {'d':2, 'e': 3}]}]}, 'a(c(e))', {'a': [{'c':[{},{'e':3}]}]}, {'a': [{'b':1, 'c': [{'d':1}, {'d':2}]}]},), ] for index, (data, _mask, expected_result, expected_result_when_negated,) in enumerate(tests, start=1):