ChangeImport: rewrite qualified references when old_name is specified

When old_name is set, ChangeImport now also handles the `import X` /
`X.name()` pattern by rewriting qualified references and managing direct
imports. Supports aliased imports (`import X as f` / `f.name()`).

Type attribution is updated on rewritten nodes (Identifier._type,
MethodInvocation._method_type).

Also fixes two pre-existing bugs:
- AddImport: insertion point now considers standalone Import nodes
- RemoveImport: UsageCollector now skips standalone Import identifiers

Author: knutwannheden

rewrite-python/rewrite/src/rewrite/python/add_import.py

@@ -287,7 +287,7 @@ def _add_import(self, cu: CompilationUnit) -> CompilationUnit:
         insert_idx = 0
         padded_stmts = list(cu.padding.statements)
         for i, padded in enumerate(padded_stmts):
-            if isinstance(padded.element, MultiImport):
+            if isinstance(padded.element, (Import, MultiImport)):
                 insert_idx = i + 1
             elif insert_idx > 0:
                 break  # Stop after we've passed the import section

rewrite-python/rewrite/src/rewrite/python/recipes/change_import.py

@@ -14,7 +14,7 @@
 
 """Recipe to change Python imports from one module/name to another."""
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace as dc_replace
 from typing import Any, Optional
 
 from rewrite import ExecutionContext, Recipe, TreeVisitor
@@ -23,15 +23,26 @@
 from rewrite.marketplace import Python
 from rewrite.recipe import option
 from rewrite.java import J
-from rewrite.java.tree import Empty, FieldAccess, Identifier, Import
+from rewrite.java.support_types import JavaType
+from rewrite.java.tree import Empty, FieldAccess, Identifier, Import, MethodInvocation
 from rewrite.python.tree import CompilationUnit, MultiImport
 from rewrite.python.visitor import PythonVisitor
 from rewrite.python.add_import import AddImportOptions, maybe_add_import
+from rewrite.python.remove_import import RemoveImportOptions, maybe_remove_import
 
 
 _Imports = [*Python, CategoryDescriptor(display_name="Imports")]
 
 
+def _create_module_type(fqn: str) -> JavaType.Class:
+    """Create a JavaType.Class for a module from its fully qualified name."""
+    class_type = JavaType.Class()
+    class_type._flags_bit_map = 0
+    class_type._fully_qualified_name = fqn
+    class_type._kind = JavaType.FullyQualified.Kind.Class
+    return class_type
+
+
 @categorize(_Imports)
 @dataclass
 class ChangeImport(Recipe):
@@ -124,48 +135,80 @@ def __init__(self):
                 super().__init__()
                 self.has_old_import = False
                 self.old_alias: Optional[str] = None
+                self.has_direct_module_import = False
+                self.module_alias: Optional[str] = None
+                self.rewrote_qualified_refs = False
+                self._new_module_type: Optional[JavaType.Class] = None
 
             def visit_compilation_unit(self, cu: CompilationUnit, p: ExecutionContext) -> J:
-                # First pass: check if the old import exists
+                # Reset state for each compilation unit
                 self.has_old_import = False
                 self.old_alias = None
+                self.has_direct_module_import = False
+                self.module_alias = None
+                self.rewrote_qualified_refs = False
 
+                # Single pass: detect old imports and direct module imports
                 for stmt in cu.statements:
                     if isinstance(stmt, Import) and not isinstance(stmt, MultiImport):
-                        alias = self._check_for_old_single_import(stmt)
-                        if alias is not None:
-                            self.has_old_import = True
-                            self.old_alias = alias if alias != "" else None
-                            break
+                        if not self.has_old_import:
+                            alias = self._check_for_old_single_import(stmt)
+                            if alias is not None:
+                                self.has_old_import = True
+                                self.old_alias = alias if alias != "" else None
+                        if old_name and not self.has_direct_module_import:
+                            name = self._get_qualid_name(stmt.qualid)
+                            if name == old_module:
+                                self.has_direct_module_import = True
+                                self.module_alias = self._get_alias_name(stmt)
                     elif isinstance(stmt, MultiImport):
-                        alias = self._check_for_old_import(stmt)
-                        if alias is not None:
-                            self.has_old_import = True
-                            self.old_alias = alias if alias != "" else None
-                            break
-
-                if not self.has_old_import:
+                        if not self.has_old_import:
+                            alias = self._check_for_old_import(stmt)
+                            if alias is not None:
+                                self.has_old_import = True
+                                self.old_alias = alias if alias != "" else None
+                        if old_name and not self.has_direct_module_import and stmt.from_ is None:
+                            for imp in stmt.names:
+                                name = self._get_qualid_name(imp.qualid)
+                                if name == old_module:
+                                    self.has_direct_module_import = True
+                                    self.module_alias = self._get_alias_name(imp)
+                                    break
+
+                if not self.has_old_import and not self.has_direct_module_import:
                     return cu
 
                 # Visit to transform imports
                 result = super().visit_compilation_unit(cu, p)
                 if not isinstance(result, CompilationUnit):
                     return result
 
-                # Schedule adding the new import
-                alias_to_use = new_alias or self.old_alias
-                if new_name:
+                # Schedule adding the new import (only for direct import changes)
+                if self.has_old_import:
+                    alias_to_use = new_alias or self.old_alias
+                    if new_name:
+                        maybe_add_import(self, AddImportOptions(
+                            module=new_module,
+                            name=new_name,
+                            alias=alias_to_use,
+                            only_if_referenced=False
+                        ))
+                    else:
+                        maybe_add_import(self, AddImportOptions(
+                            module=new_module,
+                            alias=alias_to_use,
+                            only_if_referenced=False
+                        ))
+
+                # If we rewrote qualified references, manage the direct import
+                if self.rewrote_qualified_refs:
                     maybe_add_import(self, AddImportOptions(
                         module=new_module,
-                        name=new_name,
-                        alias=alias_to_use,
+                        alias=new_alias,
                         only_if_referenced=False
                     ))
-                else:
-                    maybe_add_import(self, AddImportOptions(
-                        module=new_module,
-                        alias=alias_to_use,
-                        only_if_referenced=False
+                    maybe_remove_import(self, RemoveImportOptions(
+                        module=old_module,
                     ))
 
                 return result
@@ -196,6 +239,76 @@ def visit_multi_import(self, multi: MultiImport, p: ExecutionContext) -> Optiona
                     # import X - remove entire import
                     return self._remove_module_from_import(multi, old_module)
 
+            def visit_method_invocation(self, method: MethodInvocation, p: ExecutionContext) -> J:
+                method = super().visit_method_invocation(method, p)
+                if not old_name or not self.has_direct_module_import:
+                    return method
+                if not isinstance(method, MethodInvocation):
+                    return method
+                if not isinstance(method.select, Identifier):
+                    return method
+                if not isinstance(method.name, Identifier):
+                    return method
+
+                select_name = method.select.simple_name
+                expected_name = self.module_alias or old_module
+                if select_name != expected_name:
+                    return method
+                if method.name.simple_name != old_name:
+                    return method
+
+                self.rewrote_qualified_refs = True
+                new_select_name = new_alias or new_module
+                new_select = method.select.replace(_simple_name=new_select_name)
+                # Update type attribution on the select identifier
+                if method.select.type is not None:
+                    new_select = new_select.replace(_type=self._get_new_module_type())
+                new_padded_select = method.padding.select.replace(_element=new_select)
+                result = method.padding.replace(_select=new_padded_select)
+                if new_name and new_name != old_name:
+                    result = result.replace(_name=result.name.replace(_simple_name=new_name))
+                # Update method_type declaring type and name
+                if result.method_type is not None:
+                    result = result.replace(_method_type=dc_replace(
+                        result.method_type,
+                        _declaring_type=self._get_new_module_type(),
+                        _name=new_name or old_name,
+                    ))
+                return result
+
+            def visit_field_access(self, field_access: FieldAccess, p: ExecutionContext) -> J:
+                field_access = super().visit_field_access(field_access, p)
+                if not old_name or not self.has_direct_module_import:
+                    return field_access
+                if not isinstance(field_access, FieldAccess):
+                    return field_access
+                if not isinstance(field_access.target, Identifier):
+                    return field_access
+
+                existing_name = field_access.target.simple_name
+                expected_name = self.module_alias or old_module
+                if existing_name != expected_name:
+                    return field_access
+                if field_access.name.simple_name != old_name:
+                    return field_access
+
+                self.rewrote_qualified_refs = True
+                new_target_name = new_alias or new_module
+                new_target = field_access.target.replace(_simple_name=new_target_name)
+                # Update type attribution on the target identifier
+                if field_access.target.type is not None:
+                    new_target = new_target.replace(_type=self._get_new_module_type())
+                result = field_access.replace(_target=new_target)
+                if new_name and new_name != old_name:
+                    new_name_ident = result.name.replace(_simple_name=new_name)
+                    result = result.padding.replace(_name=result.padding.name.replace(_element=new_name_ident))
+                return result
+
+            def _get_new_module_type(self) -> JavaType.Class:
+                if self._new_module_type is None:
+                    self._new_module_type = _create_module_type(new_module)
+                return self._new_module_type
+
             def _check_for_old_single_import(self, imp: Import) -> Optional[str]:
                 """Check if a standalone J.Import matches the old import."""
                 if old_name:

rewrite-python/rewrite/src/rewrite/python/remove_import.py

@@ -111,6 +111,13 @@ def __init__(self):
                 super().__init__()
                 self.in_import = False
 
+            def visit_import(self, import_: Import, p) -> J:
+                # Don't collect identifiers from standalone import statements
+                self.in_import = True
+                result = super().visit_import(import_, p)
+                self.in_import = False
+                return result
+
             def visit_multi_import(self, multi: MultiImport, p) -> J:
                 # Don't collect identifiers from import statements
                 self.in_import = True

rewrite-python/rewrite/tests/recipes/test_change_import.py

@@ -120,3 +120,109 @@ def test_change_removes_one_name_from_multi_import(self):
                 """,
             )
         )
+
+    def test_change_qualified_method_call(self):
+        """Change: import fractions / fractions.gcd() -> import math / math.gcd()"""
+        spec = RecipeSpec(recipe=ChangeImport(
+            old_module='fractions',
+            old_name='gcd',
+            new_module='math',
+        ))
+        spec.rewrite_run(
+            python(
+                """
+                import fractions
+                result = fractions.gcd(12, 8)
+                """,
+                """
+                import math
+                result = math.gcd(12, 8)
+                """,
+            )
+        )
+
+    def test_change_aliased_qualified_method_call(self):
+        """Change: import fractions as f / f.gcd() -> import math / math.gcd()"""
+        spec = RecipeSpec(recipe=ChangeImport(
+            old_module='fractions',
+            old_name='gcd',
+            new_module='math',
+        ))
+        spec.rewrite_run(
+            python(
+                """
+                import fractions as f
+                result = f.gcd(12, 8)
+                """,
+                """
+                import math
+                result = math.gcd(12, 8)
+                """,
+            )
+        )
+
+    def test_change_qualified_ref_keeps_import_when_other_usages(self):
+        """import fractions stays when fractions.Fraction is still used."""
+        spec = RecipeSpec(recipe=ChangeImport(
+            old_module='fractions',
+            old_name='gcd',
+            new_module='math',
+        ))
+        spec.rewrite_run(
+            python(
+                """
+                import fractions
+                result = fractions.gcd(12, 8)
+                f = fractions.Fraction(1, 3)
+                """,
+                """
+                import fractions
+                import math
+                result = math.gcd(12, 8)
+                f = fractions.Fraction(1, 3)
+                """,
+            )
+        )
+
+    def test_change_qualified_field_access(self):
+        """Change: import fractions / fn = fractions.gcd -> import math / fn = math.gcd"""
+        spec = RecipeSpec(recipe=ChangeImport(
+            old_module='fractions',
+            old_name='gcd',
+            new_module='math',
+        ))
+        spec.rewrite_run(
+            python(
+                """
+                import fractions
+                fn = fractions.gcd
+                """,
+                """
+                import math
+                fn = math.gcd
+                """,
+            )
+        )
+
+    def test_change_qualified_ref_with_different_new_name(self):
+        """Qualified ref rewrite when new_name differs from old_name."""
+        spec = RecipeSpec(recipe=ChangeImport(
+            old_module='fractions',
+            old_name='gcd',
+            new_module='math',
+            new_name='greatest_common_divisor',
+        ))
+        spec.rewrite_run(
+            python(
+                """
+                import fractions
+                result = fractions.gcd(12, 8)
+                fn = fractions.gcd
+                """,
+                """
+                import math
+                result = math.greatest_common_divisor(12, 8)
+                fn = math.greatest_common_divisor
+                """,
+            )
+        )