added tests for the new with_keys option.

modified:   README.rst
	modified:   nested_lookup/nested_lookup.py
	modified:   test_nested_loopkup.py
This commit is contained in:
Russell Ballestrini 2018-04-24 12:50:00 -07:00
parent 7bfbf5ada4
commit 756ad31763
3 changed files with 79 additions and 53 deletions

View file

@ -1,30 +1,34 @@
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))
from collections import defaultdict
def _nested_lookup(key, document, wild=False, output='list'):
def nested_lookup(key, document, wild=False, with_keys=False):
"""Lookup a key in a nested document, return a list of values"""
if with_keys:
d = defaultdict(list)
for k, v in _nested_lookup(key, document, wild=wild, with_keys=with_keys):
d[k].append(v)
return d
return list(_nested_lookup(key, document, wild=wild, with_keys=with_keys))
def _nested_lookup(key, document, wild=False, with_keys=False):
"""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):
for result in _nested_lookup(key, d, wild=wild, with_keys=with_keys):
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':
if with_keys:
yield k, v
else:
yield v
elif isinstance(v, dict):
for result in _nested_lookup(key, v, wild=wild, output=output):
for result in _nested_lookup(key, v, wild=wild, with_keys=with_keys):
yield result
elif isinstance(v, list):
for d in v:
for result in _nested_lookup(key, d, wild=wild, output=output):
for result in _nested_lookup(key, d, wild=wild, with_keys=with_keys):
yield result