nested-lookup/nested_lookup/nested_lookup.py

30 lines
1.2 KiB
Python

from six import iteritems
def nested_lookup(key, document, wild=False, output='list'):
"""Lookup a key in a nested document, return a list of values"""
if output == 'dict':
return dict(_nested_lookup(key, document, wild=wild, output=output))
else:
return list(_nested_lookup(key, document, wild=wild, output=output))
def _nested_lookup(key, document, wild=False, output='list'):
"""Lookup a key in a nested document, yield a value"""
if isinstance(document, list):
for d in document:
for result in _nested_lookup(key, d, wild=wild):
yield result
if isinstance(document, dict):
for k, v in iteritems(document):
if key == k or (wild and key.lower() in k.lower()):
if output == 'dict':
yield k, v
else:
yield v
elif isinstance(v, dict):
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, output=output):
yield result