done

Author: codegen-bot

src/codegen/sdk/codebase/node_classes/node_classes.py

@@ -16,6 +16,7 @@
     from codegen.sdk.core.file import SourceFile
     from codegen.sdk.core.function import Function
     from codegen.sdk.core.import_resolution import Import
+    from codegen.sdk.core.interfaces.editable import Editable
     from codegen.sdk.core.statements.comment import Comment
     from codegen.sdk.core.symbol import Symbol
 
@@ -33,7 +34,7 @@ class NodeClasses:
     function_call_cls: type[FunctionCall]
     comment_cls: type[Comment]
     bool_conversion: dict[bool, str]
-    dynamic_import_parent_types: set[str]
+    dynamic_import_parent_types: set[type[Editable]]
     symbol_map: dict[str, type[Symbol]] = field(default_factory=dict)
     expression_map: dict[str, type[Expression]] = field(default_factory=dict)
     type_map: dict[str, type[Type] | dict[str, type[Type]]] = field(default_factory=dict)

src/codegen/sdk/codebase/node_classes/py_node_classes.py

@@ -14,7 +14,13 @@
 from codegen.sdk.core.expressions.subscript_expression import SubscriptExpression
 from codegen.sdk.core.expressions.unary_expression import UnaryExpression
 from codegen.sdk.core.expressions.unpack import Unpack
+from codegen.sdk.core.function import Function
 from codegen.sdk.core.statements.comment import Comment
+from codegen.sdk.core.statements.for_loop_statement import ForLoopStatement
+from codegen.sdk.core.statements.if_block_statement import IfBlockStatement
+from codegen.sdk.core.statements.switch_statement import SwitchStatement
+from codegen.sdk.core.statements.try_catch_statement import TryCatchStatement
+from codegen.sdk.core.statements.while_statement import WhileStatement
 from codegen.sdk.core.symbol_groups.dict import Dict
 from codegen.sdk.core.symbol_groups.list import List
 from codegen.sdk.core.symbol_groups.tuple import Tuple
@@ -29,6 +35,8 @@
 from codegen.sdk.python.expressions.string import PyString
 from codegen.sdk.python.expressions.union_type import PyUnionType
 from codegen.sdk.python.statements.import_statement import PyImportStatement
+from codegen.sdk.python.statements.match_case import PyMatchCase
+from codegen.sdk.python.statements.with_statement import WithStatement
 
 
 def parse_subscript(node: TSNode, file_node_id, ctx, parent):
@@ -110,16 +118,13 @@ def parse_subscript(node: TSNode, file_node_id, ctx, parent):
         False: "False",
     },
     dynamic_import_parent_types={
-        "function_definition",
-        "if_statement",
-        "try_statement",
-        "with_statement",
-        "else_clause",
-        "for_statement",
-        "except_clause",
-        "while_statement",
-        "match_statement",
-        "case_clause",
-        "finally_clause",
+        Function,
+        IfBlockStatement,
+        TryCatchStatement,
+        WithStatement,
+        ForLoopStatement,
+        WhileStatement,
+        SwitchStatement,
+        PyMatchCase,
     },
 )

src/codegen/sdk/codebase/node_classes/ts_node_classes.py

@@ -15,7 +15,14 @@
 from codegen.sdk.core.expressions.unary_expression import UnaryExpression
 from codegen.sdk.core.expressions.unpack import Unpack
 from codegen.sdk.core.expressions.value import Value
+from codegen.sdk.core.function import Function
 from codegen.sdk.core.statements.comment import Comment
+from codegen.sdk.core.statements.for_loop_statement import ForLoopStatement
+from codegen.sdk.core.statements.if_block_statement import IfBlockStatement
+from codegen.sdk.core.statements.switch_case import SwitchCase
+from codegen.sdk.core.statements.switch_statement import SwitchStatement
+from codegen.sdk.core.statements.try_catch_statement import TryCatchStatement
+from codegen.sdk.core.statements.while_statement import WhileStatement
 from codegen.sdk.core.symbol_groups.list import List
 from codegen.sdk.core.symbol_groups.type_parameters import TypeParameters
 from codegen.sdk.typescript.class_definition import TSClass
@@ -166,18 +173,12 @@ def parse_new(node: TSNode, *args):
         False: "false",
     },
     dynamic_import_parent_types={
-        "function_declaration",
-        "method_definition",
-        "arrow_function",
-        "if_statement",
-        "try_statement",
-        "else_clause",
-        "catch_clause",
-        "finally_clause",
-        "while_statement",
-        "for_statement",
-        "do_statement",
-        "switch_case",
-        "switch_statement",
+        Function,
+        IfBlockStatement,
+        TryCatchStatement,
+        ForLoopStatement,
+        WhileStatement,
+        SwitchStatement,
+        SwitchCase,
     },
 )

src/codegen/sdk/core/codebase.py

@@ -5,6 +5,7 @@
 import logging
 import os
 import re
+import tempfile
 from collections.abc import Generator
 from contextlib import contextmanager
 from functools import cached_property
@@ -1333,19 +1334,16 @@ def from_string(
         prog_lang = ProgrammingLanguage(language.upper()) if isinstance(language, str) else language
         filename = "test.ts" if prog_lang == ProgrammingLanguage.TYPESCRIPT else "test.py"
 
-        # Create temporary directory
-        import tempfile
+        with tempfile.TemporaryDirectory(prefix="codegen_") as tmp_dir:
+            logger.info(f"Using directory: {tmp_dir}")
 
-        tmp_dir = tempfile.mkdtemp(prefix="codegen_")
-        logger.info(f"Using directory: {tmp_dir}")
-
-        # Create codebase using factory
-        from codegen.sdk.codebase.factory.codebase_factory import CodebaseFactory
+            # Create codebase using factory
+            from codegen.sdk.codebase.factory.codebase_factory import CodebaseFactory
 
-        files = {filename: code}
-        codebase = CodebaseFactory.get_codebase_from_files(repo_path=tmp_dir, files=files, programming_language=prog_lang)
-        logger.info("Codebase initialization complete")
-        return codebase
+            files = {filename: code}
+            codebase = CodebaseFactory.get_codebase_from_files(repo_path=tmp_dir, files=files, programming_language=prog_lang)
+            logger.info("Codebase initialization complete")
+            return codebase
 
     @classmethod
     def from_files(
@@ -1411,18 +1409,15 @@ def from_files(
             prog_lang = inferred_lang
             logger.info(f"Using language: {prog_lang} ({'inferred' if language is None else 'explicit'})")
 
-        # Create temporary directory
-        import tempfile
+        with tempfile.TemporaryDirectory(prefix="codegen_") as tmp_dir:
+            logger.info(f"Using directory: {tmp_dir}")
 
-        tmp_dir = tempfile.mkdtemp(prefix="codegen_")
-        logger.info(f"Using directory: {tmp_dir}")
-
-        # Create codebase using factory
-        from codegen.sdk.codebase.factory.codebase_factory import CodebaseFactory
+            # Create codebase using factory
+            from codegen.sdk.codebase.factory.codebase_factory import CodebaseFactory
 
-        codebase = CodebaseFactory.get_codebase_from_files(repo_path=tmp_dir, files=files, programming_language=prog_lang)
-        logger.info("Codebase initialization complete")
-        return codebase
+            codebase = CodebaseFactory.get_codebase_from_files(repo_path=tmp_dir, files=files, programming_language=prog_lang)
+            logger.info("Codebase initialization complete")
+            return codebase
 
     def get_modified_symbols_in_pr(self, pr_id: int) -> tuple[str, dict[str, str], list[str]]:
         """Get all modified symbols in a pull request"""

src/codegen/sdk/core/file.py

@@ -467,15 +467,10 @@ def parse(self, ctx: CodebaseContext) -> None:
         self.code_block = self._parse_code_block(self.ts_node)
 
         self.code_block.parse()
-        self._parse_imports()
         # We need to clear the valid symbol/import names before we start resolving exports since these can be outdated.
         self.invalidate()
         sort_editables(self._nodes)
 
-    @abstractmethod
-    @commiter
-    def _parse_imports(self) -> None: ...
-
     @noapidoc
     @commiter
     def remove_internal_edges(self) -> None:

src/codegen/sdk/core/import_resolution.py

@@ -419,15 +419,7 @@ def my_function():
             bool: True if the import is dynamic (within a control flow or scope block),
             False if it's a top-level import.
         """
-        curr = self.ts_node
-
-        # always traverses upto the module level
-        while curr:
-            if curr.type in self.ctx.node_classes.dynamic_import_parent_types:
-                return True
-            curr = curr.parent
-
-        return False
+        return self.parent_of_types(self.ctx.node_classes.dynamic_import_parent_types) is not None
 
     ####################################################################################################################
     # MANIPULATIONS

src/codegen/sdk/core/interfaces/editable.py

@@ -823,7 +823,7 @@ def children_by_field_types(self, field_types: str | Iterable[str]) -> Generator
     @reader
     @noapidoc
     def child_by_field_types(self, field_types: str | Iterable[str]) -> Expression[Self] | None:
-        """Get child by field types."""
+        """Get child by fiexld types."""
         return next(self.children_by_field_types(field_types), None)
 
     @property
@@ -1097,6 +1097,14 @@ def parent_of_type(self, type: type[T]) -> T | None:
             return self.parent.parent_of_type(type)
         return None
 
+    def parent_of_types(self, types: set[type[T]]) -> T | None:
+        """Find the first ancestor of the node of the given type. Does not return itself"""
+        if self.parent and any(isinstance(self.parent, t) for t in types):
+            return self.parent
+        if self.parent is not self and self.parent is not None:
+            return self.parent.parent_of_types(types)
+        return None
+
     @reader
     def ancestors(self, type: type[T]) -> list[T]:
         """Find all ancestors of the node of the given type. Does not return itself"""

src/codegen/sdk/core/parser.py

@@ -8,6 +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
 
 if TYPE_CHECKING:
@@ -80,6 +81,42 @@ 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)
@@ -108,6 +145,7 @@ def parse_ts_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
         from codegen.sdk.typescript.statements.comment import TSComment
         from codegen.sdk.typescript.statements.for_loop_statement import TSForLoopStatement
         from codegen.sdk.typescript.statements.if_block_statement import TSIfBlockStatement
+        from codegen.sdk.typescript.statements.import_statement import TSImportStatement
         from codegen.sdk.typescript.statements.labeled_statement import TSLabeledStatement
         from codegen.sdk.typescript.statements.switch_statement import TSSwitchStatement
         from codegen.sdk.typescript.statements.try_catch_statement import TSTryCatchStatement
@@ -117,11 +155,13 @@ def parse_ts_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
 
         if node.type in self.expressions or node.type == "expression_statement":
             return [ExpressionStatement(node, file_node_id, ctx, parent, 0, expression_node=node)]
+
         for child in node.named_children:
             # =====[ Functions + Methods ]=====
             if child.type in _VALID_TYPE_NAMES:
                 statements.append(SymbolStatement(child, file_node_id, ctx, parent, len(statements)))
-
+            elif child.type == "import_statement":
+                statements.append(TSImportStatement(child, file_node_id, ctx, parent, len(statements)))
             # =====[ Classes ]=====
             elif child.type in ("class_declaration", "abstract_class_declaration"):
                 statements.append(SymbolStatement(child, file_node_id, ctx, parent, len(statements)))
@@ -132,7 +172,10 @@ def parse_ts_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
 
             # =====[ Type Alias Declarations ]=====
             elif child.type == "type_alias_declaration":
-                statements.append(SymbolStatement(child, file_node_id, ctx, parent, len(statements)))
+                if import_node := self.get_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)))
 
             # =====[ Enum Declarations ]=====
             elif child.type == "enum_declaration":
@@ -142,10 +185,10 @@ 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
+            # # =====[ 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":
@@ -167,6 +210,8 @@ 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):
+                    statements.append(TSImportStatement(import_node, file_node_id, ctx, parent, len(statements)))
                 else:
                     statements.append(
                         TSAssignmentStatement.from_assignment(
@@ -176,6 +221,10 @@ 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):
+                    statements.append(TSImportStatement(import_node, file_node_id, ctx, parent, pos=len(statements)))
+                    continue
+
                 for var in child.named_children:
                     if var.type == "string":
                         statements.append(TSComment.from_code_block(var, parent, pos=len(statements)))
@@ -185,7 +234,6 @@ def parse_ts_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
                         statements.append(ExpressionStatement(child, file_node_id, ctx, parent, pos=len(statements), expression_node=var))
             elif child.type in self.expressions:
                 statements.append(ExpressionStatement(child, file_node_id, ctx, parent, len(statements), expression_node=child))
-
             else:
                 self.log("Couldn't parse statement with type: %s", child.type)
                 statements.append(Statement.from_code_block(child, parent, pos=len(statements)))
@@ -204,6 +252,7 @@ def parse_py_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
         from codegen.sdk.python.statements.comment import PyComment
         from codegen.sdk.python.statements.for_loop_statement import PyForLoopStatement
         from codegen.sdk.python.statements.if_block_statement import PyIfBlockStatement
+        from codegen.sdk.python.statements.import_statement import PyImportStatement
         from codegen.sdk.python.statements.match_statement import PyMatchStatement
         from codegen.sdk.python.statements.pass_statement import PyPassStatement
         from codegen.sdk.python.statements.try_catch_statement import PyTryCatchStatement
@@ -237,9 +286,7 @@ def parse_py_statements(self, node: TSNode, file_node_id: NodeId, ctx: CodebaseC
 
             # =====[ Imports ] =====
             elif child.type in ["import_statement", "import_from_statement", "future_import_statement"]:
-                # statements.append(PyImportStatement(child, file_node_id, ctx, parent, len(statements)))
-                pass  # Temporarily opting to identify all imports using find_all_descendants
-
+                statements.append(PyImportStatement(child, file_node_id, ctx, parent, len(statements)))
             # =====[ Non-symbol statements ] =====
             elif child.type == "comment":
                 statements.append(PyComment.from_code_block(child, parent, pos=len(statements)))