From dfc2483e3ea5c13863b9593f84a3d293de829c9f Mon Sep 17 00:00:00 2001 From: Guillaume Date: Tue, 24 Apr 2018 09:47:47 -0700 Subject: [PATCH] added `dict` as a possible output when the user wants to keep some structure in the result, the `ouptut` option allows to return a `dict` with the matched keys as keys and corresponding values as values. --- nested_lookup/nested_lookup.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/nested_lookup/nested_lookup.py b/nested_lookup/nested_lookup.py index d5a7226..6e2f368 100644 --- a/nested_lookup/nested_lookup.py +++ b/nested_lookup/nested_lookup.py @@ -1,10 +1,15 @@ from six import iteritems -def nested_lookup(key, document, wild=False): +def nested_lookup(key, document, wild=False, output='dict'): """Lookup a key in a nested document, return a list of values""" - return list(_nested_lookup(key, document, wild=wild)) + if output == 'dict': + return dict(_nested_lookup(key, document, wild=wild, output=output)) + elif output == 'list': + return list(_nested_lookup(key, document, wild=wild, output=output)) + else: + raise ValueError('Output is of unknown type '+output) -def _nested_lookup(key, document, wild=False): +def _nested_lookup(key, document, wild=False, output='dict'): """Lookup a key in a nested document, yield a value""" if isinstance(document, list): for d in document: @@ -14,12 +19,14 @@ def _nested_lookup(key, document, wild=False): if isinstance(document, dict): for k, v in iteritems(document): if key == k or (wild and key.lower() in k.lower()): - yield v + if output == 'dict': + yield k, v + elif output == 'list': + yield v elif isinstance(v, dict): - for result in _nested_lookup(key, v, wild=wild): + for result in _nested_lookup(key, v, wild=wild, output=output): yield result elif isinstance(v, list): for d in v: - for result in _nested_lookup(key, d, wild=wild): + for result in _nested_lookup(key, d, wild=wild, output=output): yield result -