[CG-7930] new api for removing unused symbols

src/codegen/sdk/codebase/transaction_manager.py

@@ -1,3 +1,4 @@
+import math
 import time
 from collections.abc import Callable
 from pathlib import Path
@@ -289,6 +290,22 @@ def get_transactions_at_range(self, file_path: Path, start_byte: int, end_byte:
 
         return matching_transactions
 
+    def get_transaction_containing_range(self, file_path: Path, start_byte: int, end_byte: int, transaction_order: TransactionPriority | None = None) -> Transaction | None:
+        """Returns the nearest transaction that includes the range specified given the filtering criteria."""
+        if file_path not in self.queued_transactions:
+            return None
+
+        smallest_difference = math.inf
+        best_fit_transaction = None
+        for t in self.queued_transactions[file_path]:
+            if t.start_byte <= start_byte and t.end_byte >= end_byte:
+                if transaction_order is None or t.transaction_order == transaction_order:
+                    smallest_difference = min(smallest_difference, abs(t.start_byte - start_byte) + abs(t.end_byte - end_byte))
+                    if smallest_difference == 0:
+                        return t
+                    best_fit_transaction = t
+        return best_fit_transaction
+
     def _get_conflicts(self, transaction: Transaction) -> list[Transaction]:
         """Returns all transactions that overlap with the given transaction"""
         overlapping_transactions = []

src/codegen/sdk/core/file.py

@@ -943,6 +943,13 @@ def remove_unused_exports(self) -> None:
             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()
+
     ####################################################################################################################
     # MANIPULATIONS
     ####################################################################################################################

src/codegen/sdk/core/import_resolution.py

@@ -221,6 +221,17 @@ 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.
@@ -234,6 +245,16 @@ def is_wildcard_import(self) -> bool:
         """
         return self.import_type == ImportType.WILDCARD
 
+    @reader
+    def is_sideffect_import(self) -> bool:
+        # Maybe better name for this
+        """Determines if this is a sideffect.
+
+        Returns:
+            bool: True if this is a sideffect import, False otherwise
+        """
+        return self.import_type == ImportType.SIDE_EFFECT
+
     @property
     @abstractmethod
     def namespace(self) -> str | None:
@@ -663,7 +684,7 @@ def __eq__(self, other: object):
     @reader
     def remove_if_unused(self) -> None:
         if all(
-            self.transaction_manager.get_transactions_at_range(self.filepath, start_byte=usage.match.start_byte, end_byte=usage.match.end_byte, transaction_order=TransactionPriority.Remove)
+            self.transaction_manager.get_transaction_containing_range(self.file.path, start_byte=usage.match.start_byte, end_byte=usage.match.end_byte, transaction_order=TransactionPriority.Remove)
             for usage in self.usages
         ):
             self.remove()

src/codegen/sdk/core/symbol.py

@@ -5,6 +5,7 @@
 
 from rich.markup import escape
 
+from codegen.sdk.codebase.transactions import TransactionPriority
 from codegen.sdk.core.autocommit import commiter, reader, writer
 from codegen.sdk.core.dataclasses.usage import UsageKind, UsageType
 from codegen.sdk.core.detached_symbols.argument import Argument
@@ -271,6 +272,7 @@ def move_to_file(
         file: SourceFile,
         include_dependencies: bool = True,
         strategy: Literal["add_back_edge", "update_all_imports", "duplicate_dependencies"] = "update_all_imports",
+        cleanup_unused_imports: bool = True,
     ) -> None:
         """Moves the given symbol to a new file and updates its imports and references.
 
@@ -290,7 +292,91 @@ def move_to_file(
             AssertionError: If an invalid strategy is provided.
         """
         encountered_symbols = {self}
-        self._move_to_file(file, encountered_symbols, include_dependencies, strategy)
+        self._move_to_file(file, encountered_symbols, include_dependencies, strategy, cleanup_unused_imports)
+
+    @noapidoc
+    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)
+
+    @noapidoc
+    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)
+
+    @noapidoc
+    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()
 
     @noapidoc
     def _move_to_file(
@@ -299,97 +385,74 @@ def _move_to_file(
         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]:
         """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 self.transaction_manager.get_transaction_containing_range(usage.usage_symbol.file.path, usage.usage_symbol.start_byte, usage.usage_symbol.end_byte, TransactionPriority.Remove)
+            ]
+        ):
+            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()
+        from codegen.sdk.core.import_resolution import Import
+
+        self._execute_post_move_correction_strategy(file, encountered_symbols, strategy)
 
-        # ======[ 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()
+        # =====[ Remove any imports that are no longer used ]=====
+        if cleanup_unused_imports:
+            for dep in self.dependencies:
+                if strategy != "duplicate_dependencies":
+                    other_usages = [usage.usage_symbol for usage in dep.usages if usage.usage_symbol not in encountered_symbols]
+                else:
+                    other_usages = [usage.usage_symbol for usage in dep.usages]
+                if isinstance(dep, Import):
+                    dep.remove_if_unused()
+
+                    if not other_usages:
+                        # Clean up forward...
+                        dep.remove()
+                elif isinstance(dep, Symbol):
+                    usages_in_file = [
+                        symb
+                        for symb in other_usages
+                        if symb.file == self.file and not self.transaction_manager.get_transaction_containing_range(symb.file.path, symb.start_byte, symb.end_byte, TransactionPriority.Remove)
+                    ]
+                    if self.transaction_manager.get_transaction_containing_range(dep.file.path, dep.start_byte, dep.end_byte, transaction_order=TransactionPriority.Remove):
+                        if not usages_in_file and strategy != "add_back_edge":
+                            # We are going to assume there is only one such import
+                            if imp_list := [import_str for import_str in self.file._pending_imports if dep.name and dep.name in import_str]:
+                                if insert_import_list := [
+                                    transaction
+                                    for transaction in self.transaction_manager.queued_transactions[self.file.path]
+                                    if imp_list[0] and transaction.new_content and imp_list[0] in transaction.new_content and transaction.transaction_order == TransactionPriority.Insert
+                                ]:
+                                    self.transaction_manager.queued_transactions[self.file.path].remove(insert_import_list[0])
+                                    self.file._pending_imports.remove(imp_list[0])
 
     @property
     @reader