diff --git a/src/codegen/extensions/tools/search_files_by_name.py b/src/codegen/extensions/tools/search_files_by_name.py index b44f6da85..829376cf3 100644 --- a/src/codegen/extensions/tools/search_files_by_name.py +++ b/src/codegen/extensions/tools/search_files_by_name.py @@ -1,7 +1,7 @@ import math import shutil import subprocess -from typing import ClassVar, Optional +from typing import ClassVar from pydantic import Field @@ -88,13 +88,13 @@ def search_files_by_name( if files_per_page == math.inf: files_per_page = total_files total_pages = 1 - else: + else: total_pages = (total_files + files_per_page - 1) // files_per_page if total_files > 0 else 1 - - + + # Ensure page is within valid range page = min(page, total_pages) - + # Get paginated results start_idx = (page - 1) * files_per_page end_idx = start_idx + files_per_page diff --git a/src/codegen/sdk/core/file.py b/src/codegen/sdk/core/file.py index e5af34ef9..73b381f6d 100644 --- a/src/codegen/sdk/core/file.py +++ b/src/codegen/sdk/core/file.py @@ -946,9 +946,7 @@ def remove_unused_exports(self) -> None: def remove_unused_imports(self) -> None: # Process each import statement for import_stmt in self.imports: - # Don't remove imports we can't be sure about - if import_stmt.usage_is_ascertainable(): - import_stmt.remove_if_unused() + import_stmt.remove_if_unused() #################################################################################################################### # MANIPULATIONS diff --git a/src/codegen/sdk/core/import_resolution.py b/src/codegen/sdk/core/import_resolution.py index 1fb17df50..f4f7b7b33 100644 --- a/src/codegen/sdk/core/import_resolution.py +++ b/src/codegen/sdk/core/import_resolution.py @@ -220,17 +220,6 @@ def is_symbol_import(self) -> bool: """ return not self.is_module_import() - @reader - def usage_is_ascertainable(self) -> bool: - """Returns True if we can determine for sure whether the import is unused or not. - - Returns: - bool: True if the usage can be ascertained for the import, False otherwise. - """ - if self.is_wildcard_import() or self.is_sideffect_import(): - return False - return True - @reader def is_wildcard_import(self) -> bool: """Returns True if the import symbol is a wildcard import. @@ -691,7 +680,7 @@ def remove_if_unused(self, force: bool = False) -> bool: bool: True if removed, False if not """ if all(usage.match.get_transaction_if_pending_removal() for usage in self.usages): - if not force and not self.usage_is_ascertainable(): + if not force and (self.is_wildcard_import() or self.is_sideffect_import()): return False self.remove() return True diff --git a/src/codegen/sdk/core/symbol.py b/src/codegen/sdk/core/symbol.py index bce4a91e1..3987624d0 100644 --- a/src/codegen/sdk/core/symbol.py +++ b/src/codegen/sdk/core/symbol.py @@ -267,6 +267,7 @@ def insert_before(self, new_src: str, fix_indentation: bool = False, newline: bo return first_node.insert_before(new_src, fix_indentation, newline, priority, dedupe) return super().insert_before(new_src, fix_indentation, newline, priority, dedupe) + @noapidoc def _post_move_import_cleanup(self, encountered_symbols, strategy): # =====[ Remove any imports that are no longer used ]===== from codegen.sdk.core.import_resolution import Import @@ -293,6 +294,88 @@ def _post_move_import_cleanup(self, encountered_symbols, strategy): self.transaction_manager.queued_transactions[self.file.path].remove(insert_import_list[0]) self.file._pending_imports.remove(imp_list[0]) + def _move_dependencies( + self, file: SourceFile, encountered_symbols: set[Symbol | Import], strategy: Literal["add_back_edge", "update_all_imports", "duplicate_dependencies"] = "update_all_imports" + ): + from codegen.sdk.core.import_resolution import Import + + for dep in self.dependencies: + if dep in encountered_symbols: + continue + + # =====[ Symbols - move over ]===== + if isinstance(dep, Symbol) and dep.is_top_level: + encountered_symbols.add(dep) + dep._move_to_file( + file=file, + encountered_symbols=encountered_symbols, + include_dependencies=True, + strategy=strategy, + ) + + # =====[ Imports - copy over ]===== + elif isinstance(dep, Import): + if dep.imported_symbol: + file.add_import(imp=dep.imported_symbol, alias=dep.alias.source, import_type=dep.import_type) + else: + file.add_import(imp=dep.source) + + def _update_dependencies_on_move(self, file: SourceFile): + from codegen.sdk.core.import_resolution import Import + + for dep in self.dependencies: + # =====[ Symbols - add back edge ]===== + if isinstance(dep, Symbol) and dep.is_top_level: + file.add_import(imp=dep, alias=dep.name, import_type=ImportType.NAMED_EXPORT, is_type_import=False) + elif isinstance(dep, Import): + if dep.imported_symbol: + file.add_import(imp=dep.imported_symbol, alias=dep.alias.source) + else: + file.add_import(imp=dep.source) + + + def _execute_post_move_correction_strategy( + self, + file: SourceFile, + encountered_symbols: set[Symbol | Import], + strategy: Literal["add_back_edge", "update_all_imports", "duplicate_dependencies"], + ): + from codegen.sdk.core.import_resolution import Import + + import_line = self.get_import_string(module=file.import_module_name) + is_used_in_file = any(usage.file == self.file and usage.node_type == NodeType.SYMBOL and usage not in encountered_symbols for usage in self.symbol_usages) + match strategy: + case "duplicate_dependencies": + # If not used in the original file. or if not imported from elsewhere, we can just remove the original symbol + if not is_used_in_file and not any(usage.kind is UsageKind.IMPORTED and usage.usage_symbol not in encountered_symbols for usage in self.usages): + self.remove() + case "add_back_edge": + if is_used_in_file or any(usage.kind is UsageKind.IMPORTED and usage.usage_symbol not in encountered_symbols for usage in self.usages): + self.file.add_import(imp=import_line) + + # Delete the original symbol + self.remove() + case "update_all_imports": + for usage in self.usages: + if isinstance(usage.usage_symbol, Import) and usage.usage_symbol.file != file: + # Add updated import + usage.usage_symbol.file.add_import(import_line) + usage.usage_symbol.remove() + elif usage.usage_type == UsageType.CHAINED: + # Update all previous usages of import * to the new import name + if usage.match and "." + self.name in usage.match: + if isinstance(usage.match, FunctionCall) and self.name in usage.match.get_name(): + usage.match.get_name().edit(self.name) + if isinstance(usage.match, ChainedAttribute): + usage.match.edit(self.name) + usage.usage_symbol.file.add_import(imp=import_line) + + # Add the import to the original file + if is_used_in_file: + self.file.add_import(imp=import_line) + # Delete the original symbol + self.remove() + def move_to_file( self, file: SourceFile, @@ -330,96 +413,41 @@ def _move_to_file( cleanup_unused_imports: bool = True, ) -> tuple[NodeId, NodeId]: """Helper recursive function for `move_to_file`""" - from codegen.sdk.core.import_resolution import Import - # =====[ Arg checking ]===== if file == self.file: return file.file_node_id, self.node_id + + # This removes any imports in the file we're moving the symbol to if imp := file.get_import(self.name): encountered_symbols.add(imp) imp.remove() + # This handles the new file + # =====[ Move over dependencies recursively ]===== if include_dependencies: - # =====[ Move over dependencies recursively ]===== - for dep in self.dependencies: - if dep in encountered_symbols: - continue - - # =====[ Symbols - move over ]===== - if isinstance(dep, Symbol) and dep.is_top_level: - encountered_symbols.add(dep) - dep._move_to_file( - file=file, - encountered_symbols=encountered_symbols, - include_dependencies=include_dependencies, - strategy=strategy, - ) - - # =====[ Imports - copy over ]===== - elif isinstance(dep, Import): - if dep.imported_symbol: - file.add_import(imp=dep.imported_symbol, alias=dep.alias.source) - else: - file.add_import(imp=dep.source) + self._move_dependencies(file, encountered_symbols, strategy) else: - for dep in self.dependencies: - # =====[ Symbols - add back edge ]===== - if isinstance(dep, Symbol) and dep.is_top_level: - file.add_import(imp=dep, alias=dep.name, import_type=ImportType.NAMED_EXPORT, is_type_import=False) - elif isinstance(dep, Import): - if dep.imported_symbol: - file.add_import(imp=dep.imported_symbol, alias=dep.alias.source) - else: - file.add_import(imp=dep.source) + self._update_dependencies_on_move(file) # =====[ Make a new symbol in the new file ]===== - file.add_symbol(self) - import_line = self.get_import_string(module=file.import_module_name) - + should_export = False + if hasattr(self, "is_exported") and ( + self.is_exported + or [ + usage + for usage in self.usages + if usage.usage_symbol not in encountered_symbols + and not usage.usage_symbol.get_transaction_if_pending_removal() + ] + ): + should_export = True + file.add_symbol(self, should_export=should_export) # =====[ Checks if symbol is used in original file ]===== # Takes into account that it's dependencies will be moved - is_used_in_file = any( - usage.file == self.file and usage.node_type == NodeType.SYMBOL and usage not in encountered_symbols and (usage.start_byte < self.start_byte or usage.end_byte > self.end_byte) # HACK - for usage in self.symbol_usages - ) - - # ======[ Strategy: Duplicate Dependencies ]===== - if strategy == "duplicate_dependencies": - # If not used in the original file. or if not imported from elsewhere, we can just remove the original symbol - if not is_used_in_file and not any(usage.kind is UsageKind.IMPORTED and usage.usage_symbol not in encountered_symbols for usage in self.usages): - self.remove() - # ======[ Strategy: Add Back Edge ]===== - # Here, we will add a "back edge" to the old file importing the symbol - elif strategy == "add_back_edge": - if is_used_in_file or any(usage.kind is UsageKind.IMPORTED and usage.usage_symbol not in encountered_symbols for usage in self.usages): - self.file.add_import(imp=import_line) - # Delete the original symbol - self.remove() - - # ======[ Strategy: Update All Imports ]===== - # Update the imports in all the files which use this symbol to get it from the new file now - elif strategy == "update_all_imports": - for usage in self.usages: - if isinstance(usage.usage_symbol, Import) and usage.usage_symbol.file != file: - # Add updated import - usage.usage_symbol.file.add_import(import_line) - usage.usage_symbol.remove() - elif usage.usage_type == UsageType.CHAINED: - # Update all previous usages of import * to the new import name - if usage.match and "." + self.name in usage.match: - if isinstance(usage.match, FunctionCall) and self.name in usage.match.get_name(): - usage.match.get_name().edit(self.name) - if isinstance(usage.match, ChainedAttribute): - usage.match.edit(self.name) - usage.usage_symbol.file.add_import(imp=import_line) - - # Add the import to the original file - if is_used_in_file: - self.file.add_import(imp=import_line) - # Delete the original symbol - self.remove() + self._execute_post_move_correction_strategy(file, encountered_symbols, strategy) + # =====[ Remove any imports that are no longer used ]===== if cleanup_unused_imports: self._post_move_import_cleanup(encountered_symbols, strategy) diff --git a/src/codegen/sdk/typescript/symbol.py b/src/codegen/sdk/typescript/symbol.py index 903fd8806..a80b06da9 100644 --- a/src/codegen/sdk/typescript/symbol.py +++ b/src/codegen/sdk/typescript/symbol.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal, Self, Unpack +from typing import TYPE_CHECKING, Literal, Self, Unpack, override from codegen.sdk.core.assignment import Assignment from codegen.sdk.core.autocommit import reader, writer @@ -27,7 +27,6 @@ from codegen.sdk.core.file import SourceFile from codegen.sdk.core.import_resolution import Import from codegen.sdk.core.interfaces.editable import Editable - from codegen.sdk.core.node_id_factory import NodeId from codegen.sdk.typescript.detached_symbols.code_block import TSCodeBlock from codegen.sdk.typescript.interfaces.has_block import TSHasBlock @@ -254,133 +253,75 @@ def has_semicolon(self) -> bool: """ return self.semicolon_node is not None - @noapidoc - def _move_to_file( - self, - file: SourceFile, - encountered_symbols: set[Symbol | Import], - include_dependencies: bool = True, - strategy: Literal["add_back_edge", "update_all_imports", "duplicate_dependencies"] = "update_all_imports", - cleanup_unused_imports: bool = True, - ) -> tuple[NodeId, NodeId]: - # TODO: Prevent creation of import loops (!) - raise a ValueError and make the agent fix it - # =====[ Arg checking ]===== - if file == self.file: - return file.file_node_id, self.node_id - - if imp := file.get_import(self.name): - encountered_symbols.add(imp) - imp.remove() - - # =====[ Move over dependencies recursively ]===== - if include_dependencies: - try: - for dep in self.dependencies: - if dep in encountered_symbols: - continue - - # =====[ Symbols - move over ]===== - elif isinstance(dep, TSSymbol): - if dep.is_top_level: - encountered_symbols.add(dep) - dep._move_to_file(file, encountered_symbols=encountered_symbols, include_dependencies=True, strategy=strategy) - - # =====[ Imports - copy over ]===== - elif isinstance(dep, TSImport): - if dep.imported_symbol: - file.add_import(dep.imported_symbol, alias=dep.alias.source, import_type=dep.import_type) - else: - file.add_import(dep.source) - - else: - msg = f"Unknown dependency type {type(dep)}" - raise ValueError(msg) - except Exception as e: - print(f"Failed to move dependencies of {self.name}: {e}") - else: - try: - for dep in self.dependencies: - if isinstance(dep, Assignment): - msg = "Assignment not implemented yet" - raise NotImplementedError(msg) - - # =====[ Symbols - move over ]===== - elif isinstance(dep, Symbol) and dep.is_top_level: - file.add_import(imp=dep, alias=dep.name, import_type=ImportType.NAMED_EXPORT, is_type_import=isinstance(dep, TypeAlias)) - - if not dep.is_exported: - dep.file.add_export_to_symbol(dep) - pass - - # =====[ Imports - copy over ]===== - elif isinstance(dep, TSImport): - if dep.imported_symbol: - file.add_import(dep.imported_symbol, alias=dep.alias.source, import_type=dep.import_type, is_type_import=dep.is_type_import()) - else: - file.add_import(dep.source) - - except Exception as e: - print(f"Failed to move dependencies of {self.name}: {e}") - - # =====[ Make a new symbol in the new file ]===== - # This will update all edges etc. - should_export = False - - if self.is_exported or [usage for usage in self.usages if usage.usage_symbol not in encountered_symbols and not usage.usage_symbol.get_transaction_if_pending_removal()]: - should_export = True - - file.add_symbol(self, should_export=should_export) - import_line = self.get_import_string(module=file.import_module_name) - - # =====[ Checks if symbol is used in original file ]===== - # Takes into account that it's dependencies will be moved - is_used_in_file = any(usage.file == self.file and usage.node_type == NodeType.SYMBOL and usage not in encountered_symbols for usage in self.symbol_usages) - - # ======[ Strategy: Duplicate Dependencies ]===== - if strategy == "duplicate_dependencies": - # If not used in the original file. or if not imported from elsewhere, we can just remove the original symbol - is_used_in_file = any(usage.file == self.file and usage.node_type == NodeType.SYMBOL for usage in self.symbol_usages) - if not is_used_in_file and not any(usage.kind is UsageKind.IMPORTED and usage.usage_symbol not in encountered_symbols for usage in self.usages): - self.remove() - - # ======[ Strategy: Add Back Edge ]===== - # Here, we will add a "back edge" to the old file importing the self - elif strategy == "add_back_edge": - if is_used_in_file: - back_edge_line = import_line - if self.is_exported: - back_edge_line = back_edge_line.replace("import", "export") - self.file.add_import(back_edge_line) - elif self.is_exported: - module_name = file.name - self.file.add_import(f"export {{ {self.name} }} from '{module_name}'") - # Delete the original symbol - self.remove() - - # ======[ Strategy: Update All Imports ]===== - # Update the imports in all the files which use this symbol to get it from the new file now - elif strategy == "update_all_imports": - for usage in self.usages: - if isinstance(usage.usage_symbol, TSImport) and usage.usage_symbol.file != file: - # Add updated import - usage.usage_symbol.file.add_import(import_line) - usage.usage_symbol.remove() - elif usage.usage_type == UsageType.CHAINED: - # Update all previous usages of import * to the new import name - if usage.match and "." + self.name in usage.match: - if isinstance(usage.match, FunctionCall) and self.name in usage.match.get_name(): - usage.match.get_name().edit(self.name) - if isinstance(usage.match, ChainedAttribute): - usage.match.edit(self.name) - usage.usage_symbol.file.add_import(imp=import_line) - - # Add the import to the original file - if is_used_in_file: - self.file.add_import(imp=import_line) - # Delete the original symbol - self.remove() - if cleanup_unused_imports: - self._post_move_import_cleanup(encountered_symbols, strategy) + @override + def _update_dependencies_on_move(self,file:SourceFile): + for dep in self.dependencies: + if isinstance(dep, Assignment): + msg = "Assignment not implemented yet" + raise NotImplementedError(msg) + + # =====[ Symbols - move over ]===== + elif isinstance(dep, TSSymbol) and dep.is_top_level: + file.add_import(imp=dep, alias=dep.name, import_type=ImportType.NAMED_EXPORT, is_type_import=isinstance(dep, TypeAlias)) + + if not dep.is_exported: + dep.file.add_export_to_symbol(dep) + pass + # =====[ Imports - copy over ]===== + elif isinstance(dep, TSImport): + if dep.imported_symbol: + file.add_import(dep.imported_symbol, alias=dep.alias.source, import_type=dep.import_type, is_type_import=dep.is_type_import()) + else: + file.add_import(dep.source) + + @override + def _execute_post_move_correction_strategy(self, + file:SourceFile, + encountered_symbols: set[Symbol | Import], + strategy: Literal["add_back_edge", "update_all_imports", "duplicate_dependencies"], + ): + import_line = self.get_import_string(module=file.import_module_name) + # =====[ Checks if symbol is used in original file ]===== + # Takes into account that it's dependencies will be moved + is_used_in_file = any(usage.file == self.file and usage.node_type == NodeType.SYMBOL and usage not in encountered_symbols for usage in self.symbol_usages) + + match strategy: + case "duplicate_dependencies": + # If not used in the original file. or if not imported from elsewhere, we can just remove the original symbol + is_used_in_file = any(usage.file == self.file and usage.node_type == NodeType.SYMBOL for usage in self.symbol_usages) + if not is_used_in_file and not any(usage.kind is UsageKind.IMPORTED and usage.usage_symbol not in encountered_symbols for usage in self.usages): + self.remove() + case "add_back_edge": + if is_used_in_file: + back_edge_line = import_line + if self.is_exported: + back_edge_line=back_edge_line.replace('import','export') + self.file.add_import(back_edge_line) + elif self.is_exported: + module_name = file.name + self.file.add_import(f"export {{ {self.name} }} from '{module_name}'") + # Delete the original symbol + self.remove() + case "update_all_imports": + for usage in self.usages: + if isinstance(usage.usage_symbol, TSImport) and usage.usage_symbol.file != file: + # Add updated import + usage.usage_symbol.file.add_import(import_line) + usage.usage_symbol.remove() + elif usage.usage_type == UsageType.CHAINED: + # Update all previous usages of import * to the new import name + if usage.match and "." + self.name in usage.match: + if isinstance(usage.match, FunctionCall) and self.name in usage.match.get_name(): + usage.match.get_name().edit(self.name) + if isinstance(usage.match, ChainedAttribute): + usage.match.edit(self.name) + usage.usage_symbol.file.add_import(import_line) + + if is_used_in_file: + self.file.add_import(import_line) + # Delete the original symbol + self.remove() + def _convert_proptype_to_typescript(self, prop_type: Editable, param: Parameter | None, level: int) -> str: """Converts a PropType definition to its TypeScript equivalent.""" @@ -388,6 +329,7 @@ def _convert_proptype_to_typescript(self, prop_type: Editable, param: Parameter type_map = {"string": "string", "number": "number", "bool": "boolean", "object": "object", "array": "any[]", "func": "CallableFunction"} if prop_type.source in type_map: return type_map[prop_type.source] + if isinstance(prop_type, ChainedAttribute): if prop_type.attribute.source == "node": return "T" diff --git a/tests/unit/codegen/sdk/python/function/test_function_move_to_file.py b/tests/unit/codegen/sdk/python/function/test_function_move_to_file.py index a4c29dcdc..533ea90fa 100644 --- a/tests/unit/codegen/sdk/python/function/test_function_move_to_file.py +++ b/tests/unit/codegen/sdk/python/function/test_function_move_to_file.py @@ -498,6 +498,43 @@ def baz(): FILE_4_CONTENT = """ from file2 import bar +def bla(): + return bar() + 1 +""" + + # ========== [ AFTER ] ========== + # language=python + EXPECTED_FILE_1_CONTENT = """ +def external_dep(): + return 42 +""" + + # language=python + EXPECTED_FILE_2_CONTENT = """ +from file3 import bar +def foo(): + return foo_dep() + 1 + +def foo_dep(): + return 24 + +def bar(): + return external_dep() + bar_dep() + +def bar_dep(): + return 2 +""" + + # language=python + FILE_3_CONTENT = """ +from file2 import bar + +def baz(): + return bar() + 1 +""" + FILE_4_CONTENT = """ +from file2 import bar + def bla(): return bar() + 1 """ @@ -1583,13 +1620,13 @@ def test_move_to_file_decorators(tmpdir) -> None: @test_decorator.foo() def test_func(): - pass - """ + pass""" FILE_2_CONTENT = "" EXPECTED_FILE_1_CONTENT = "" - EXPECTED_FILE_2_CONTENT = """from test.foo import TEST + EXPECTED_FILE_2_CONTENT =\ +"""from test.foo import TEST test_decorator = TEST() @@ -1615,8 +1652,8 @@ def test_func(): file1 = codebase.get_file("file1.py") file2 = codebase.get_file("file2.py") - assert file1.source == EXPECTED_FILE_1_CONTENT - assert file2.source == EXPECTED_FILE_2_CONTENT + assert file1.source == EXPECTED_FILE_1_CONTENT.strip('/n') + assert file2.source == EXPECTED_FILE_2_CONTENT.strip('/n') def test_move_to_file_multiple_same_transaction(tmpdir) -> None: