diff --git a/README.rst b/README.rst index 10c4c44..7a29c39 100644 --- a/README.rst +++ b/README.rst @@ -30,8 +30,8 @@ or install from source using:: cd nested-lookup pip install . -tutorial -======== +quick tutorial +============== .. code-block:: python @@ -42,14 +42,20 @@ tutorial >>> print(nested_lookup('taco', document)) [42, 69] +longer tutorial +=============== -wild -======== +You may control the libraries behavior by passing some optional arguments. -We also have a `wild` mode that treats the given `key` as a case insensitive -substring of all the keys in the document and returns any values which match. +wild (defaults to `False`): + if `wild` is `True`, treat the given `key` as a case insensitive + substring when performing lookups. -For example: +return_keys (defaults to `False`): + if `with_keys` is `True`, return a dictionary of all matched keys + and a list of values. + +For example, given the following document: .. code-block:: python @@ -61,9 +67,14 @@ For example: 'other' : { 'secondary_email' : 'test2@example.com', 'EMAIL_RECOVERY' : 'test3@example.com', + 'email_address' : 'test4@example.com', }, }, +We could act `wild` and find all the email addresses like this: + +.. code-block:: python + results = nested_lookup( key = 'mail', document = my_document, @@ -71,43 +82,31 @@ For example: ) print(results) - ['test1@example.com', 'test2@example.com', 'test3@example.com'] - - -output -======== - -There are two `output` modes: - -* `list`: the function returns a list of values corresponding to the matched keys. -* `dict`: the function returns a `dict` with the matched keys as keys and their corresponding values as values. - -For example: .. code-block:: python - from nested_lookup import nested_lookup + ['test1@example.com', 'test4@example.com', 'test2@example.com', 'test3@example.com'] - my_document = { - 'name' : 'Russell Ballestrini', - 'email_address' : 'test1@example.com', - 'other' : { - 'secondary_email' : 'test2@example.com', - 'EMAIL_RECOVERY' : 'test3@example.com', - }, - }, +Additionally, if you also needed the matched key names, you could do this: + +.. code-block:: python results = nested_lookup( key = 'mail', document = my_document, wild = True, - output = 'dict' + with_keys = True, ) print(results) - {'email_address': 'test1@example.com', - 'secondary_email': 'test2@example.com', - 'EMAIL_RECOVERY': 'test3@example.com'} + +.. code-block:: python + + { + 'email_address': ['test1@example.com', 'test4@example.com'], + 'secondary_email': ['test2@example.com'], + 'EMAIL_RECOVERY': ['test3@example.com'] + } misc diff --git a/nested_lookup/nested_lookup.py b/nested_lookup/nested_lookup.py index 645b657..78ec193 100644 --- a/nested_lookup/nested_lookup.py +++ b/nested_lookup/nested_lookup.py @@ -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 diff --git a/test_nested_loopkup.py b/test_nested_loopkup.py index 5a2af8d..178335d 100644 --- a/test_nested_loopkup.py +++ b/test_nested_loopkup.py @@ -6,6 +6,15 @@ class TestNestedLookup(TestCase): def setUp(self): self.subject_dict = {'a':1,'b':{'d':100},'c':{'d':200}} + self.subject_dict2 = { + 'name' : 'Russell Ballestrini', + 'email_address' : 'test1@example.com', + 'other' : { + 'secondary_email' : 'test2@example.com', + 'EMAIL_RECOVERY' : 'test3@example.com', + 'email_address' : 'test4@example.com', + }, + } def test_nested_lookup(self): results = nested_lookup('d', self.subject_dict) @@ -37,18 +46,32 @@ class TestNestedLookup(TestCase): def test_wild_nested_lookup(self): results = nested_lookup( - key = 'mail', - document = { - 'name' : 'Russell Ballestrini', - 'email_address' : 'test1@example.com', - 'other' : { - 'secondary_email' : 'test2@example.com', - 'EMAIL_RECOVERY' : 'test3@example.com', - }, - }, + key = 'mail', + document = self.subject_dict2 wild = True, ) - self.assertEqual(3, len(results)) + self.assertEqual(4, len(results)) self.assertIn('test1@example.com', results) self.assertIn('test2@example.com', results) self.assertIn('test3@example.com', results) + + def test_wild_with_keys_nested_lookup(self): + matches = nested_lookup( + key = 'mail', + document = self.subject_dict2, + wild = True, + with_keys = True, + ) + self.assertEqual(3, len(matches)) + self.assertIn('email_address', matches) + self.assertIn('secondary_email', matches) + self.assertIn('EMAIL_RECOVERY', matches) + self.assertSetEqual({'test1@example.com','test4@example.com'}, set(matches['email_address'])) + self.assertIn('test2@example.com', matches['secondary_email']) + + def test_nested_lookup_with_keys(self): + matches = nested_lookup('d', self.subject_dict, with_keys=True) + self.assertIn('d', matches) + self.assertEqual(2, len(matches['d'])) + self.assertSetEqual({100,200}, set(matches['d'])) +