added tests for the new with_keys option.

modified:   README.rst
	modified:   nested_lookup/nested_lookup.py
	modified:   test_nested_loopkup.py

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

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

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']))
+