[CG-7930] new api for removing unused symbols (#855)

continuation of #855

---------

Co-authored-by: tomcodegen <kucar.tomislav@gmail.com>
Co-authored-by: tomcodgen <191515280+tomcodgen@users.noreply.github.com>

Author: 3 people

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

@@ -5,7 +5,6 @@
 from typing import TYPE_CHECKING, ClassVar, Generic, Literal, Self, TypeVar, override
 
 from codegen.sdk.codebase.resolution_stack import ResolutionStack
-from codegen.sdk.codebase.transactions import TransactionPriority
 from codegen.sdk.core.autocommit import commiter, reader, remover, writer
 from codegen.sdk.core.dataclasses.usage import UsageKind
 from codegen.sdk.core.expressions.name import Name
@@ -221,6 +220,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 +244,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:
@@ -661,12 +681,21 @@ def __eq__(self, other: object):
 
     @noapidoc
     @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)
-            for usage in self.usages
-        ):
+    def remove_if_unused(self, force: bool = False) -> bool:
+        """Removes import if it is not being used. Considers current transaction removals.
+
+        Args:
+            force (bool, optional): If true removes the import even if we cannot ascertain the usage for sure. Defaults to False.
+
+        Returns:
+            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():
+                return False
             self.remove()
+            return True
+        return False
 
     @noapidoc
     @reader

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

@@ -10,7 +10,7 @@
 from rich.pretty import Pretty
 
 from codegen.sdk.codebase.span import Span
-from codegen.sdk.codebase.transactions import EditTransaction, InsertTransaction, RemoveTransaction, TransactionPriority
+from codegen.sdk.codebase.transactions import EditTransaction, InsertTransaction, RemoveTransaction, Transaction, TransactionPriority
 from codegen.sdk.core.autocommit import commiter, reader, remover, repr_func, writer
 from codegen.sdk.core.placeholder.placeholder import Placeholder
 from codegen.sdk.extensions.utils import get_all_identifiers
@@ -1156,6 +1156,15 @@ def parent_class(self) -> Class | None:
 
         return self.parent_of_type(Class)
 
+    @noapidoc
+    def get_transaction_if_pending_removal(self) -> Transaction | None:
+        """Checks if this editable is being removed by some transaction and if so returns it.
+
+        Returns:
+            Transaction|None: The transaction removing the editable
+        """
+        return self.transaction_manager.get_transaction_containing_range(self.file.path, self.start_byte, self.end_byte, TransactionPriority.Remove)
+
     def _get_ast_children(self) -> list[tuple[str | None, AST]]:
         children = []
         names = {}

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
@@ -266,11 +267,38 @@ 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)
 
+    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
+
+        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()
+
+            elif isinstance(dep, Symbol):
+                usages_in_file = [symb for symb in other_usages if symb.file == self.file and not symb.get_transaction_if_pending_removal()]
+                if dep.get_transaction_if_pending_removal():
+                    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])
+
     def move_to_file(
         self,
         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 +318,7 @@ 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_to_file(
@@ -299,6 +327,7 @@ 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
@@ -391,6 +420,9 @@ def _move_to_file(
             # Delete the original symbol
             self.remove()
 
+        if cleanup_unused_imports:
+            self._post_move_import_cleanup(encountered_symbols, strategy)
+
     @property
     @reader
     @noapidoc

src/codegen/sdk/typescript/symbol.py

@@ -261,12 +261,17 @@ 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]:
         # 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:
@@ -319,7 +324,12 @@ def _move_to_file(
 
         # =====[ Make a new symbol in the new file ]=====
         # This will update all edges etc.
-        file.add_symbol(self)
+        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 ]=====
@@ -329,16 +339,18 @@ def _move_to_file(
         # ======[ 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:
-                self.file.add_import(import_line)
+                back_edge_line = import_line
                 if self.is_exported:
-                    self.file.add_import(f"export {{ {self.name} }}")
+                    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}'")
@@ -349,23 +361,26 @@ def _move_to_file(
         # 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):
+                if isinstance(usage.usage_symbol, TSImport) and usage.usage_symbol.file != file:
                     # Add updated import
-                    if usage.usage_symbol.resolved_symbol is not None and usage.usage_symbol.resolved_symbol.node_type == NodeType.SYMBOL and usage.usage_symbol.resolved_symbol == self:
-                        usage.usage_symbol.file.add_import(import_line)
-                        usage.usage_symbol.remove()
+                    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):
+                        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)
+                        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(import_line)
+                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)
 
     def _convert_proptype_to_typescript(self, prop_type: Editable, param: Parameter | None, level: int) -> str:
         """Converts a PropType definition to its TypeScript equivalent."""