Merge pull request #12 from Agent-Hellboy/issue-10

Fix #10

Author: Agent-Hellboy

src/trace_dkey/__init__.py

@@ -27,20 +27,23 @@ def __trace_key(dict_key: Any, dict_value: Any, key: Any, path: List[Union[str,
     if dict_key_value == key:
         paths = [path]
     else:
-        # branch to iterate over list of keys
+        # branch to iterate over list of keys 
+        # this branch will iterate over keys like {'a':1,'b':2,'c':3}
         if (type(dict_value) == _ast.Dict and len(dict_value.keys) > 0):
             for i in range(len(dict_value.keys)):
                 # find array of paths inside the ith branch
                 branch_paths = __trace_key(dict_value.keys[i], dict_value.values[i], key, path)
                 paths = [*paths, *branch_paths]
 
-        # branch to iterate over dict of list  
+        # branch to iterate over dict of list 
+        # this branch will iterate over keys like {'a':1,'b':2,'c':[{'d':'e']}} but skip if c value is list of other types ['d','e']
         if (type(dict_value) == _ast.List):
             for elt in dict_value.elts:
-                for i in range(len(elt.keys)):
-                    # find array of paths inside the ith branch
-                    branch_paths = __trace_key(elt.keys[i], elt.values[i], key, path)
-                    paths = [*paths, *branch_paths]
+                if type(elt) == _ast.Dict:
+                    for i in range(len(elt.keys)):
+                        # find array of paths inside the ith branch
+                        branch_paths = __trace_key(elt.keys[i], elt.values[i], key, path)
+                        paths = [*paths, *branch_paths]
     return paths
 
 

tests/test_trace_dkey.py

@@ -19,7 +19,9 @@ def test_trace_fixture():
     fixture4 = trace(dict4, 'c')
     dict5={'a':'b','c':{'d':[{'e':'f'}]}}
     fixture5 = trace(dict5,'e')
-    return fixture1, fixture2, fixture3, fixture4, fixture5
+    dict6={'a':'b','c':{'d':['e','f']}}
+    fixture6 = trace(dict6,'e')
+    return fixture1, fixture2, fixture3, fixture4, fixture5, fixture6
 
 
 def test_when_key_is_string(test_trace_fixture):
@@ -35,4 +37,7 @@ def test_when_multiple_keys_found(test_trace_fixture):
     assert test_trace_fixture[3] == [['a', 'b', 'c'], ['a', 'c']]
 
 def test_when_keys_found_in_list(test_trace_fixture):
-    assert test_trace_fixture[4] == [['c', 'd','e']]
+    assert test_trace_fixture[4] == [['c', 'd','e']]
+
+def test_when_keys_not_found_in_list(test_trace_fixture):
+    assert test_trace_fixture[5] == []