[CG-7930] moving symbols refactor

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

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

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

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)