feat: support deeper navigation and reference

src/ast_selector/models.py

@@ -30,10 +30,10 @@ def append(self, selector_group: SelectorGroup, val: ast.AST) -> None:
         ref = self._build_reference(selector_group)
         self.reference_table[ref].append(val)
 
-    def get_ref_idx(self, ref: str) -> Optional[int]:
+    def get_reference(self, ref: str) -> Optional[List[ast.AST]]:
         k = list(reversed(self.reference_table.keys()))
         if ref in k:
-            return k.index(ref)
+            return self.reference_table[ref]
 
         return None
 
@@ -119,10 +119,12 @@ def _find_nodes(self, branches: Union[Iterator[ast.AST], ast.AST]) -> Generator[
             branches = iter([branches])
 
         for tree in branches:
-            parent = tree
             for node in ast.walk(tree):
-                node.parent = parent  # type: ignore
-                parent = node
+                node.parent = getattr(node, "parent", tree)  # type: ignore
+
+                for child in ast.iter_child_nodes(node):
+                    child.parent = node  # type: ignore
+
                 if isinstance(node, self.element_type):
                     if self._matches(node):
                         yield node
@@ -205,20 +207,23 @@ class ReferenceSelector(ElementSelector):
     def __post_init__(self) -> None:
         self.element_type = ast.AST
 
+    def _walk_parent(self, node: ast.AST) -> Generator[ast.AST, None, None]:
+        cur_node = node
+        while cur_node.parent and cur_node != cur_node.parent:  # type: ignore
+            yield cur_node.parent  # type: ignore
+            cur_node = cur_node.parent  # type: ignore
+
     def _find_nodes(self, branches: Union[Iterator[ast.AST], ast.AST]) -> Generator[ast.AST, None, None]:
         if not isinstance(branches, Iterator):
             branches = iter([branches])
 
         tree = list(branches)  # Force generators to build reference table until this point
-        parent_level = self.navigation.get_ref_idx(self.query)
-        if parent_level:
-            for node in tree:
-                selected = node
-                for _ in range(parent_level):
-                    selected = selected.parent  # type: ignore
-
-                if self._matches(selected):
-                    yield selected
+        references = self.navigation.get_reference(self.query) or []
+
+        for ref in references:
+            if any((node for node in tree if ref in self._walk_parent(node))):
+                if self._matches(ref):
+                    yield ref
 
 
 @dataclass

src/tests/ast_selector_test.py

@@ -200,6 +200,16 @@ def test_drill_list_ast_from_reference():
     assert all(isinstance(x, ast.AST) for x in found)
 
 
+def test_long_navigation_get_back_to_original_reference():
+    tree = read_sample("log_object")
+    query = "FunctionDef Expr[value is Call].value[func is Attribute].func[attr = exception] $FunctionDef"
+
+    selector = AstSelector(query, tree)
+    found = selector.first()
+
+    assert isinstance(found, ast.FunctionDef)
+
+
 # TODO: Support array (e.g. FunctionDef.body.0)
 # TODO: Support direct children (e.g. FunctionDef > FunctionDef using ast.iter_children instead of ast.walk)
 # TODO: Support deeper references (e.g. FunctionDef[1] when FunctionDef(not this) FunctionDef(this))