fixing errors

Author: tawsifkamal

src/codegen/sdk/core/parser.py

@@ -8,8 +8,7 @@
 from codegen.sdk.core.expressions.placeholder_type import PlaceholderType
 from codegen.sdk.core.expressions.value import Value
 from codegen.sdk.core.statements.symbol_statement import SymbolStatement
-from codegen.sdk.extensions.utils import find_all_descendants, find_first_descendant
-from codegen.sdk.utils import find_first_function_descendant
+from codegen.sdk.utils import find_first_function_descendant, find_import_node
 
 if TYPE_CHECKING:
     from tree_sitter import Node as TSNode
@@ -81,42 +80,6 @@ def parse_expression(self, node: TSNode | None, file_node_id: NodeId, ctx: Codeb
                 ret.children
         return ret
 
-    def get_import_node(self, node: TSNode) -> TSNode | None:
-        """Get the import node from a node that may contain an import.
-        Returns None if the node does not contain an import.
-
-        Returns:
-            TSNode | None: The import_statement or call_expression node if it's an import, None otherwise
-        """
-        # Static imports
-        if node.type == "import_statement":
-            return node
-
-        # Dynamic imports and requires can be either:
-        # 1. Inside expression_statement -> call_expression
-        # 2. Direct call_expression
-
-        # we only parse imports inside expressions and variable declarations
-        call_expression = find_first_descendant(node, ["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"])
-            if descendants:
-                import_node = descendants[-1]
-            else:
-                # this means this is NOT a dynamic import()
-                return None
-        else:
-            import_node = call_expression
-
-        # thus we only consider the deepest one
-        if import_node:
-            function = import_node.child_by_field_name("function")
-            if function and (function.type == "import" or (function.type == "identifier" and function.text.decode("utf-8") == "require")):
-                return import_node
-
-        return None
-
     def log_unparsed(self, node: TSNode) -> None:
         if self._should_log and node.is_named and node.type not in self._uncovered_nodes:
             self._uncovered_nodes.add(node.type)
@@ -172,7 +135,7 @@ def parse_ts_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
 
             # =====[ Type Alias Declarations ]=====
             elif child.type == "type_alias_declaration":
-                if import_node := self.get_import_node(child):
+                if import_node := find_import_node(child):
                     statements.append(TSImportStatement(import_node, file_node_id, ctx, parent, len(statements)))
                 else:
                     statements.append(SymbolStatement(child, file_node_id, ctx, parent, len(statements)))
@@ -185,11 +148,6 @@ def parse_ts_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
             elif child.type == "export_statement" or child.text.decode("utf-8") == "export *;":
                 statements.append(ExportStatement(child, file_node_id, ctx, parent, len(statements)))
 
-            # # =====[ Imports ] =====
-            # elif child.type == "import_statement":
-            #     # statements.append(TSImportStatement(child, file_node_id, ctx, parent, len(statements)))
-            #     pass  # Temporarily opting to identify all imports using find_all_descendants
-
             # =====[ Non-symbol statements ] =====
             elif child.type == "comment":
                 statements.append(TSComment.from_code_block(child, parent, pos=len(statements)))
@@ -210,7 +168,7 @@ def parse_ts_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
             elif child.type in ["lexical_declaration", "variable_declaration"]:
                 if function_node := find_first_function_descendant(child):
                     statements.append(SymbolStatement(child, file_node_id, ctx, parent, len(statements), function_node))
-                elif import_node := self.get_import_node(child):
+                elif import_node := find_import_node(child):
                     statements.append(TSImportStatement(import_node, file_node_id, ctx, parent, len(statements)))
                 else:
                     statements.append(
@@ -221,7 +179,7 @@ def parse_ts_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
             elif child.type in ["public_field_definition", "property_signature", "enum_assignment"]:
                 statements.append(TSAttribute(child, file_node_id, ctx, parent, pos=len(statements)))
             elif child.type == "expression_statement":
-                if import_node := self.get_import_node(child):
+                if import_node := find_import_node(child):
                     statements.append(TSImportStatement(import_node, file_node_id, ctx, parent, pos=len(statements)))
                     continue
 

src/codegen/sdk/utils.py

@@ -87,6 +87,43 @@ def find_first_function_descendant(node: TSNode) -> TSNode:
     return find_first_descendant(node=node, type_names=type_names, max_depth=2)
 
 
+def find_import_node(node: TSNode) -> TSNode | None:
+    """Get the import node from a node that may contain an import.
+    Returns None if the node does not contain an import.
+
+    Returns:
+        TSNode | None: The import_statement or call_expression node if it's an import, None otherwise
+    """
+    # Static imports
+    if node.type == "import_statement":
+        return node
+
+    # Dynamic imports and requires can be either:
+    # 1. Inside expression_statement -> call_expression
+    # 2. Direct call_expression
+
+    # we only parse imports inside expressions and variable declarations
+
+    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"])
+        if descendants:
+            import_node = descendants[-1]
+        else:
+            # this means this is NOT a dynamic import()
+            return None
+    else:
+        import_node = find_first_descendant(node, ["call_expression"])
+
+    # thus we only consider the deepest one
+    if import_node:
+        function = import_node.child_by_field_name("function")
+        if function and (function.type == "import" or (function.type == "identifier" and function.text.decode("utf-8") == "require")):
+            return import_node
+
+    return None
+
+
 def find_index(target: TSNode, siblings: list[TSNode]) -> int:
     """Returns the index of the target node in the list of siblings, or -1 if not found. Recursive implementation."""
     if target in siblings:

tests/unit/codegen/sdk/codebase/session/test_codebase_from_files.py

@@ -11,7 +11,7 @@ def test_from_files_python():
     assert len(codebase.files) == 2
     assert any(f.filepath.endswith("main.py") for f in codebase.files)
     assert any(f.filepath.endswith("utils.py") for f in codebase.files)
-    assert any("from utils import add" in f.content for f in codebase.files)
+    assert any("from utils import add" in f.source for f in codebase.files)
 
 
 def test_from_files_typescript():
@@ -22,7 +22,7 @@ def test_from_files_typescript():
     assert len(codebase.files) == 2
     assert any(f.filepath.endswith("index.ts") for f in codebase.files)
     assert any(f.filepath.endswith("utils.ts") for f in codebase.files)
-    assert any("import { add }" in f.content for f in codebase.files)
+    assert any("import { add }" in f.source for f in codebase.files)
 
 
 def test_from_files_empty():

tests/unit/codegen/sdk/codebase/session/test_codebase_from_string.py

@@ -13,7 +13,7 @@ def hello():
     codebase = Codebase.from_string(code, language="python")
     assert len(codebase.files) == 1
     assert codebase.files[0].filepath.endswith("test.py")
-    assert "def hello" in codebase.files[0].content
+    assert "def hello" in codebase.files[0].source
 
 
 def test_from_string_typescript():
@@ -26,7 +26,7 @@ def test_from_string_typescript():
     codebase = Codebase.from_string(code, language="typescript")
     assert len(codebase.files) == 1
     assert codebase.files[0].filepath.endswith("test.ts")
-    assert "function hello" in codebase.files[0].content
+    assert "function hello" in codebase.files[0].source
 
 
 def test_from_string_with_enum():
@@ -42,14 +42,14 @@ def test_from_string_invalid_syntax():
     code = "this is not valid python"
     codebase = Codebase.from_string(code, language="python")
     assert len(codebase.files) == 1
-    assert codebase.files[0].content == code
+    assert codebase.files[0].source == code
 
 
 def test_from_string_empty():
     """Test creating a codebase from empty string"""
     codebase = Codebase.from_string("", language="python")
     assert len(codebase.files) == 1
-    assert codebase.files[0].content == ""
+    assert codebase.files[0].source == ""
 
 
 def test_from_string_missing_language():