[CG-7930] feat: remove unused imports after moving symbol & new api for removing unused symbols

src/codegen/sdk/core/symbol.py

@@ -39,7 +39,7 @@
 
 
 @apidoc
 class Symbol(Usable[Statement["CodeBlock[Parent, ...]"]], Generic[Parent, TCodeBlock]):
     """Abstract representation of a Symbol in a Codebase. A Symbol is a top-level entity in a file, e.g. a Function, Class, GlobalVariable, etc.
 
     Attributes:
@@ -55,13 +55,13 @@
         ts_node: TSNode,
         file_id: NodeId,
         G: CodebaseGraph,
         parent: Statement[CodeBlock[Parent, ...]],
         name_node: TSNode | None = None,
         name_node_type: type[Name] = DefinedName,
     ) -> None:
         super().__init__(ts_node, file_id, G, parent)
         name_node = self._get_name_node(ts_node) if name_node is None else name_node
         self._name_node = self._parse_expression(name_node, default=name_node_type)
         from codegen.sdk.core.interfaces.has_block import HasBlock
 
         if isinstance(self, HasBlock):
@@ -73,7 +73,7 @@
     def __rich_repr__(self) -> rich.repr.Result:
         yield escape(self.filepath) + "::" + (self.full_name if self.full_name else "<no name>")
 
     __rich_repr__.angular = ANGULAR_STYLE
 
     @property
     @noapidoc
@@ -85,7 +85,7 @@
         if parent == self.file or isinstance(parent, Export):
             # Top level symbol
             return self
         return parent
 
     @staticmethod
     @noapidoc
@@ -271,6 +271,7 @@
         file: SourceFile,
         include_dependencies: bool = True,
         strategy: Literal["add_back_edge", "update_all_imports", "duplicate_dependencies"] = "update_all_imports",
+        remove_unused_imports: bool = True,
     ) -> None:
         """Moves the given symbol to a new file and updates its imports and references.
 
@@ -282,6 +283,7 @@
             strategy (str): The strategy to use for updating imports. Can be either 'add_back_edge' or 'update_all_imports'. Defaults to 'update_all_imports'.
                 - 'add_back_edge': Moves the symbol and adds an import in the original file
                 - 'update_all_imports': Updates all imports and usages of the symbol to reference the new file
+            remove_unused_imports (bool): If True, removes any imports in the original file that become unused after moving the symbol. Defaults to True.
 
         Returns:
             None
@@ -290,25 +292,30 @@
             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, remove_unused_imports)
 
     @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",
+        remove_unused_imports: bool = True,
     ) -> tuple[NodeId, NodeId]:
         """Helper recursive function for `move_to_file`"""
         from codegen.sdk.core.import_resolution import Import
 
+        # Track original file and imports used by this symbol before moving
+        symbol_imports = set()
+
         # =====[ 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()
+            if remove_unused_imports:
+                imp.remove()
 
         if include_dependencies:
             # =====[ Move over dependencies recursively ]=====
@@ -329,7 +336,7 @@
                 # =====[ Imports - copy over ]=====
                 elif isinstance(dep, Import):
                     if dep.imported_symbol:
                         file.add_symbol_import(dep.imported_symbol, alias=dep.alias.source)
                     else:
                         file.add_import_from_import_string(dep.source)
         else:
@@ -391,6 +398,40 @@
             # Delete the original symbol
             self.remove()
 
+        # After moving a symbol (function or class) out of a file, if there are imports that are now unused because that was the only thing using them, remove those as well
+        if remove_unused_imports:
+            # Get all imports that were used by the moved symbol
+            for dep in self.dependencies:
+                if isinstance(dep, Import):
+                    symbol_imports.add(dep)
+
+            # Check each import that was used by the moved symbol
+            for import_symbol in symbol_imports:
+                try:
+                    # Try to access any property - if the import was removed this will fail
+                    _ = import_symbol.file
+                except (AttributeError, ReferenceError):
+                    # Skip if import was already removed
+                    continue
+
+                # Check if import is still used by any remaining symbols
+                still_used = False
+                for usage in import_symbol.usages:
+                    # Skip usages from the moved symbol
+                    if usage.usage_symbol == self:
+                        continue
+
+                    # Skip usages from symbols we moved
+                    if usage.usage_symbol in encountered_symbols:
+                        continue
+
+                    still_used = True
+                    break
+
+                # Remove import if it's no longer used
+                if not still_used:
+                    import_symbol.remove()
+
     @property
     @reader
     @noapidoc

src/codegen/sdk/python/file.py

@@ -174,10 +174,38 @@
         else:
             self.insert_before(import_string, priority=1)
 
-    @noapidoc
+    def remove_unused_imports(self) -> None:
+        """Removes unused imports from the file.
+
+        Handles different Python import styles:
+        - Single imports (import x)
+        - From imports (from y import z)
+        - Multi-imports (from y import (a, b as c))
+        - Wildcard imports (from x import *)
+        - Type imports (from typing import List)
+        - Future imports (from __future__ import annotations)
+
+        Preserves:
+        - Comments and whitespace where possible
+        - Future imports even if unused
+        - Type hints and annotations
+        """
+        # Process each import statement
+        for import_stmt in self.imports:
+            # Always preserve __future__ and star imports since we can't track their usage
+            if import_stmt.is_future_import or import_stmt.is_wildcard_import():
+                continue
+
+            import_stmt.remove_if_unused()
+
+        self.G.commit_transactions()
+
     def remove_unused_exports(self) -> None:
-        """Removes unused exports from the file. NO-OP for python"""
-        pass
+        """Removes unused exports from the file.
+        In Python this is equivalent to removing unused imports since Python doesn't have
+        explicit export statements. Calls remove_unused_imports() internally.
+        """
+        self.remove_unused_imports()
 
     @cached_property
     @noapidoc

src/codegen/sdk/python/import_resolution.py

@@ -284,6 +284,55 @@ def get_import_string(
         else:
             return f"from {import_module} import {self.name}"
 
+    @property
+    def module_name(self) -> str:
+        """Gets the module name for this import.
+
+        For 'import x' returns 'x'
+        For 'from x import y' returns 'x'
+        For 'from .x import y' returns '.x'
+
+        Returns:
+            str: The module name for this import.
+        """
+        if self.ts_node.type == "import_from_statement":
+            module_node = self.ts_node.child_by_field_name("module_name")
+            return module_node.text.decode("utf-8") if module_node else ""
+        return self.ts_node.child_by_field_name("name").text.decode("utf-8")
+
+    def is_from_import(self) -> bool:
+        """Determines if this is a from-style import statement.
+
+        Checks if the import uses 'from' syntax (e.g., 'from module import symbol')
+        rather than direct import syntax (e.g., 'import module').
+
+        Returns True for imports like:
+        - from x import y
+        - from .x import y
+        - from x import (a, b, c)
+        - from x import *
+
+        Returns False for:
+        - import x
+        - import x as y
+
+        Returns:
+            bool: True if this is a from-style import, False otherwise.
+        """
+        return self.import_type in [ImportType.NAMED_EXPORT, ImportType.WILDCARD]
+
+    @property
+    def is_future_import(self) -> bool:
+        """Determines if this is a __future__ import.
+
+        Returns True for imports like:
+        - from __future__ import annotations
+
+        Returns:
+            bool: True if this is a __future__ import, False otherwise
+        """
+        return self.ts_node.type == "future_import_statement"
+
 
 class PyExternalImportResolver(ExternalImportResolver):
     def __init__(self, from_alias: str, to_context: CodebaseGraph) -> None:

src/codegen/sdk/typescript/file.py

@@ -6,7 +6,7 @@
 from codegen.sdk.core.autocommit import commiter, mover, reader, writer
 from codegen.sdk.core.file import SourceFile
 from codegen.sdk.core.interfaces.exportable import Exportable
-from codegen.sdk.enums import ImportType, NodeType, ProgrammingLanguage, SymbolType
+from codegen.sdk.enums import ImportType, ProgrammingLanguage, SymbolType
 from codegen.sdk.extensions.sort import sort_editables
 from codegen.sdk.extensions.utils import cached_property
 from codegen.sdk.typescript.assignment import TSAssignment
@@ -238,63 +238,6 @@
                 if function.type == "import" or (function.type == "identifier" and function.text.decode("utf-8") == "require"):
                     TSImportStatement(import_node, self.node_id, self.G, self.code_block, 0)
 
-    @writer
-    def remove_unused_exports(self) -> None:
-        """Removes unused exports from the file.
-
-        Analyzes all exports in the file and removes any that are not used. An export is considered unused if it has no direct
-        symbol usages and no re-exports that are used elsewhere in the codebase.
-
-        When removing unused exports, the method also cleans up any related unused imports. For default exports, it removes
-        the 'export default' keyword, and for named exports, it removes the 'export' keyword or the entire export statement.
-
-        Args:
-            None
-
-        Returns:
-            None
-        """
-        for export in self.exports:
-            symbol_export_unused = True
-            symbols_to_remove = []
-
-            exported_symbol = export.resolved_symbol
-            for export_usage in export.symbol_usages:
-                if export_usage.node_type == NodeType.IMPORT or (export_usage.node_type == NodeType.EXPORT and export_usage.resolved_symbol != exported_symbol):
-                    # If the import has no usages then we can add the import to the list of symbols to remove
-                    reexport_usages = export_usage.symbol_usages
-                    if len(reexport_usages) == 0:
-                        symbols_to_remove.append(export_usage)
-                        break
-
-                    # If any of the import's usages are valid symbol usages, export is used.
-                    if any(usage.node_type == NodeType.SYMBOL for usage in reexport_usages):
-                        symbol_export_unused = False
-                        break
-
-                    symbols_to_remove.append(export_usage)
-
-                elif export_usage.node_type == NodeType.SYMBOL:
-                    symbol_export_unused = False
-                    break
-
-            # export is not used, remove it
-            if symbol_export_unused:
-                # remove the unused imports
-                for imp in symbols_to_remove:
-                    imp.remove()
-
-                if exported_symbol == exported_symbol.export.declared_symbol:
-                    # change this to be more robust
-                    if exported_symbol.source.startswith("export default "):
-                        exported_symbol.replace("export default ", "")
-                    else:
-                        exported_symbol.replace("export ", "")
-                else:
-                    exported_symbol.export.remove()
-                if exported_symbol.export != export:
-                    export.remove()
-
     @noapidoc
     def _get_export_data(self, relative_path: str, export_type: str = "EXPORT") -> tuple[tuple[str, str], dict[str, callable]]:
         quoted_paths = (f"'{relative_path}'", f'"{relative_path}"')
@@ -394,11 +337,27 @@
     def valid_import_names(self) -> dict[str, Symbol | TSImport]:
         """Returns a dict mapping name => Symbol (or import) in this file that can be imported from another file"""
         valid_export_names = {}
+
+        # Handle default exports
         if len(self.default_exports) == 1:
             valid_export_names["default"] = self.default_exports[0]
+
+        # Handle named exports and their aliases
         for export in self.exports:
             for name, dest in export.names:
+                # Track both original name and alias if present
                 valid_export_names[name] = dest
+                if hasattr(dest, "alias") and dest.alias:
+                    valid_export_names[dest.alias] = dest
+
+        # # Handle imports and their aliases
+        # for import_stmt in self.imports:
+        #     for name, symbol in import_stmt.imported_symbols.items():
+        #         valid_export_names[name] = symbol
+        #         # Also track the alias if present
+        #         if hasattr(symbol, "alias") and symbol.alias:
+        #             valid_export_names[symbol.alias] = symbol
+
         return valid_export_names
 
     ####################################################################################################################
@@ -445,3 +404,84 @@
             TSNamespace | None: The namespace with the specified name if found, None otherwise.
         """
         return next((x for x in self.symbols if isinstance(x, TSNamespace) and x.name == name), None)
+
+    @writer
+    def remove_unused_imports(self) -> None:
+        """Removes unused imports from the file.
+
+        Handles different TypeScript import styles:
+        - Single imports (import x from 'y')
+        - Named imports (import { x } from 'y')
+        - Multi-imports (import { a, b as c } from 'y')
+        - Type imports (import type { X } from 'y')
+        - Side effect imports (import 'y')
+        - Wildcard imports (import * as x from 'y')
+
+        Preserves:
+        - Comments and whitespace where possible
+        - Side effect imports (e.g., CSS imports)
+        - Type imports used in type annotations
+        """
+        # Process each import statement
+        for import_stmt in self.imports:
+            # Always preserve side effect imports since we can't track their usage
+            if import_stmt.import_type == ImportType.SIDE_EFFECT:
+                continue
+
+            # Check if all imports in this statement are unused
+            import_stmt.remove_if_unused()
+
+        self.G.commit_transactions()
+
+    def _is_export_used(self, export: TSExport) -> bool:
+        # Get all symbol usages
+        usages = export.symbol_usages()
+
+        # If there are any usages, the export is used
+        if usages:
+            return True
+
+        # Check if this is a re-export that's used elsewhere
+        if export.is_reexport():
+            # Get the original symbol
+            original = export.resolved_symbol
+            if original:
+                # Check usages of the original symbol
+                return bool(original.symbol_usages())
+
+        return False
+
+    @writer
+    def remove_unused_exports(self) -> None:
+        """Removes unused exports from the file.
+
+        Handles different TypeScript export styles:
+        - Default exports (export default x)
+        - Named exports (export function x, export const x)
+        - Re-exports (export { x } from 'y')
+        - Type exports (export type X, export interface X)
+
+        Preserves:
+        - Type exports (these may be used in type positions)
+        - Default exports (these are often used dynamically)
+        - Exports used by other files through imports
+        - Exports used within the same file
+        """
+        exports_to_remove = []
+
+        for export in self.exports:
+            # Skip type exports
+            if export.is_type_export() or export.is_default_export():
+                continue
+
+            # Check if export is used
+            if self._is_export_used(export):
+                continue
+
+            exports_to_remove.append(export)
+
+        # Remove unused exports
+        for export in exports_to_remove:
+            export.remove()
+
+        self.G.commit_transactions()