diff --git a/nested_lookup/nested_lookup.py b/nested_lookup/nested_lookup.py index 71ac8e4..0ae066d 100644 --- a/nested_lookup/nested_lookup.py +++ b/nested_lookup/nested_lookup.py @@ -54,21 +54,17 @@ def get_all_keys(dictionary): """ result_list = [] - def recrusion(dictionary): - for key, value in iteritems(dictionary): - if isinstance(value, dict): - result_list.append(key) - recrusion(dictionary=value) - elif isinstance(value, list): - result_list.append(key) - for list_items in value: - # Make sure the items inside the list is iterable - if hasattr(list_items, 'items'): - recrusion(dictionary=list_items) - else: + def recrusion(document): + if isinstance(document, list): + for list_items in document: + recrusion(document=list_items) + elif isinstance(document, dict): + for key, value in iteritems(document): result_list.append(key) + recrusion(document=value) + return - recrusion(dictionary=dictionary) + recrusion(document=dictionary) return result_list diff --git a/test_nested_lookup.py b/test_nested_lookup.py index 5f44f8a..718a795 100644 --- a/test_nested_lookup.py +++ b/test_nested_lookup.py @@ -193,6 +193,16 @@ class TestGetAllKeys(TestCase): }] }] } + self.sample5 = [{ + "listings": [{ + "name": "title", + "postcode": "postcode", + "full_address": "fulladdress", + "city": "city", + "lat": "latitude", + "lng": "longitude" + }] + }] def test_sample_data1(self): result = get_all_keys(self.sample1) @@ -242,6 +252,16 @@ class TestGetAllKeys(TestCase): for key in keys_to_verify: self.assertIn(key, result) + def test_sample_data5(self): + result = get_all_keys(self.sample5) + self.assertEqual(7, len(result)) + keys_to_verify = [ + 'listings', 'name', 'postcode', 'full_address', 'city', + 'lat', 'lng' + ] + for key in keys_to_verify: + self.assertIn(key, result) + class TestGetOccurrence(TestCase): def setUp(self):