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

ruff.toml

@@ -157,8 +157,11 @@ extend-generics = [
   "codegen.sdk.python.expressions.named_type.PyNamedType",
   "codegen.sdk.python.expressions.string.PyString",
   "codegen.sdk.python.expressions.union_type.PyUnionType",
+  "codegen.sdk.python.file.remove_unused_imports",
+  "codegen.sdk.python.file.remove_unused_exports",
   "codegen.sdk.python.file.PyFile",
   "codegen.sdk.python.function.PyFunction",
+  "codegen.sdk.python.import_resolution.is_from_import",
   "codegen.sdk.python.import_resolution.PyImport",
   "codegen.sdk.python.interfaces.has_block.PyHasBlock",
   "codegen.sdk.python.placeholder.placeholder_return_type.PyReturnTypePlaceholder",

src/codegen/sdk/core/symbol.py

@@ -264,7 +264,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)
 
-    def move_to_file(self, file: SourceFile, include_dependencies: bool = True, strategy: str = "update_all_imports") -> None:
+    def move_to_file(self, file: SourceFile, include_dependencies: bool = True, strategy: str = "update_all_imports", remove_unused_imports: bool = True) -> None:
         """Moves the given symbol to a new file and updates its imports and references.
 
         This method moves a symbol to a new file and updates all references to that symbol throughout the codebase. The way imports are handled can be controlled via the strategy parameter.
@@ -275,6 +275,7 @@ def move_to_file(self, file: SourceFile, include_dependencies: bool = True, stra
             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
@@ -283,19 +284,25 @@ def move_to_file(self, file: SourceFile, include_dependencies: bool = True, stra
             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: str = "update_all_imports") -> tuple[NodeId, NodeId]:
+    def _move_to_file(
+        self, file: SourceFile, encountered_symbols: set[Symbol | Import], include_dependencies: bool = True, strategy: str = "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 ]=====
@@ -368,6 +375,40 @@ def _move_to_file(self, file: SourceFile, encountered_symbols: set[Symbol | Impo
         # =====[ 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

@@ -170,10 +170,73 @@ def add_import_from_import_string(self, import_string: str) -> None:
         else:
             self.insert_before(import_string, priority=1)
 
-    @noapidoc
+    @py_apidoc
+    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
+        """
+        # Track processed imports to avoid duplicates
+        processed_imports = set()
+
+        # Group imports by module for more efficient processing
+        module_imports = {}
+
+        # First pass - group imports by module
+        for import_stmt in self.imports:
+            if import_stmt in processed_imports:
+                continue
+
+            # Always preserve __future__ and star imports since we can't track their usage
+            print(f"import_stmt: {import_stmt}", import_stmt.is_future_import, import_stmt.is_star_import)
+            if import_stmt.is_future_import or import_stmt.is_star_import:
+                continue
+
+            module = import_stmt.module_name
+            if module not in module_imports:
+                module_imports[module] = []
+                print(f"Adding {module} to module_imports")
+            module_imports[module].append(import_stmt)
+
+        print(f"module_imports: {module_imports}")
+        # Second pass - process each module's imports
+        for module, imports in module_imports.items():
+            # Skip if any import from this module is used
+            if any(imp.usages for imp in imports):
+                # Remove individual unused imports if it's a from-style import
+                if len(imports) > 1 and imports[0].is_from_import():
+                    for imp in imports:
+                        if not imp.usages and imp not in processed_imports:
+                            processed_imports.add(imp)
+                            imp.remove()
+                continue
+
+            # If no imports from module are used, remove them all
+            for imp in imports:
+                if imp not in processed_imports:
+                    processed_imports.add(imp)
+                    imp.remove()
+
+        self.G.commit_transactions()
+
+    @py_apidoc
     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,70 @@ 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")
+
+    @py_apidoc
+    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.ts_node.type == "import_from_statement"
+
+    @property
+    def is_star_import(self) -> bool:
+        """Determines if this is a star import (from x import *).
+
+        Returns:
+            bool: True if this is a star import, False otherwise
+        """
+        if self.ts_node.type != "import_from_statement":
+            return False
+
+        # Look for wildcard_import node among children
+        wildcard_import = next((node for node in self.ts_node.children if node.type == "wildcard_import"), None)
+        return wildcard_import is not None
+
+    @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

@@ -389,11 +389,27 @@ def get_import_string(self, alias: str | None = None, module: str | None = None,
     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
 
     ####################################################################################################################
@@ -440,3 +456,53 @@ def get_namespace(self, name: str) -> TSNamespace | None:
             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, moved_symbol_names: set[str] | None = None) -> None:
+        """Removes unused imports from the file.
+
+        Args:
+            moved_symbol_names: Optional set of symbol names that were moved to another file
+        """
+        for import_statement in self.import_statements:
+            # Track which symbols in this import statement are still used
+            used_symbols = []
+            removed_symbols = []
+
+            for import_symbol in import_statement.imports:
+                # Skip side effect imports
+                if import_symbol.import_type == ImportType.SIDE_EFFECT:
+                    continue
+
+                symbol_name = import_symbol.alias.source if import_symbol.alias else import_symbol.name
+
+                # Check if this import is still used in the file
+                is_used = False
+                for usage in import_symbol.usages:
+                    # Skip usages from moved symbols if provided
+                    if moved_symbol_names and usage.usage_symbol and usage.usage_symbol.name in moved_symbol_names:
+                        continue
+                    is_used = True
+                    break
+
+                if is_used:
+                    used_symbols.append(import_symbol)
+                else:
+                    removed_symbols.append(import_symbol)
+
+            if not used_symbols and removed_symbols:
+                # If no symbols are used, remove the entire import statement
+                import_statement.remove()
+            elif removed_symbols and used_symbols:
+                # If some symbols are used but others aren't, update the import statement
+                new_imports = []
+                for symbol in used_symbols:
+                    if symbol.alias:
+                        new_imports.append(f"{symbol.name} as {symbol.alias.source}")
+                    else:
+                        new_imports.append(symbol.name)
+
+                module = import_statement.module.source
+                type_prefix = "type " if import_statement.is_type_import else ""
+                new_statement = f"import {type_prefix}{{ {', '.join(new_imports)} }} from {module};"
+                import_statement.source = new_statement