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.
This commit is contained in:
Guillaume 2018-04-24 09:47:47 -07:00
parent 1b7685db3f
commit dfc2483e3e

View file

@ -1,10 +1,15 @@
from six import iteritems 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""" """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""" """Lookup a key in a nested document, yield a value"""
if isinstance(document, list): if isinstance(document, list):
for d in document: for d in document:
@ -14,12 +19,14 @@ def _nested_lookup(key, document, wild=False):
if isinstance(document, dict): if isinstance(document, dict):
for k, v in iteritems(document): for k, v in iteritems(document):
if key == k or (wild and key.lower() in k.lower()): 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): 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 yield result
elif isinstance(v, list): elif isinstance(v, list):
for d in v: 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 yield result