done

Author: tawsifkamal

src/codegen/sdk/extensions/utils.pyi

@@ -13,6 +13,7 @@ def find_all_descendants(
     type_names: Iterable[str] | str,
     max_depth: int | None = None,
     nested: bool = True,
+    stop_at_first: str | None = None,
 ) -> list[TSNode]: ...
 def find_line_start_and_end_nodes(node: TSNode) -> list[tuple[TSNode, TSNode]]:
     """Returns a list of tuples of the start and end nodes of each line in the node"""

src/codegen/sdk/extensions/utils.pyx

@@ -31,7 +31,7 @@ def get_all_identifiers(node: TSNode) -> list[TSNode]:
     return sorted(dict.fromkeys(identifiers), key=lambda x: x.start_byte)
 
 
-def find_all_descendants(node: TSNode, type_names: Iterable[str] | str, max_depth: int | None = None, nested: bool = True) -> list[TSNode]:
+def find_all_descendants(node: TSNode, type_names: Iterable[str] | str, max_depth: int | None = None, nested: bool = True, stop_at_first: str | None = None) -> list[TSNode]:
     if isinstance(type_names, str):
         type_names = [type_names]
     descendants = []
@@ -45,6 +45,9 @@ def find_all_descendants(node: TSNode, type_names: Iterable[str] | str, max_dept
             if not nested and current_node != node:
                 return
 
+        if stop_at_first and current_node.type == stop_at_first:
+            return
+
         for child in current_node.children:
             traverse(child, depth + 1)
 

src/codegen/sdk/utils.py

@@ -104,9 +104,11 @@ def find_import_node(node: TSNode) -> TSNode | None:
 
     # we only parse imports inside expressions and variable declarations
 
+    # import_nodes = [_node for _node in find_all_descendants(node, ["call_expression", "statement_block"], nested=False) if _node.type == "call_expression"]
+
     if member_expression := find_first_descendant(node, ["member_expression"]):
         # there may be multiple call expressions (for cases such as import(a).then(module => module).then(module => module)
-        descendants = find_all_descendants(member_expression, ["call_expression"])
+        descendants = find_all_descendants(member_expression, ["call_expression"], stop_at_first="statement_block")
         if descendants:
             import_node = descendants[-1]
         else: