Python: Add type attribution tests, use moduleName for module-qualified FQNs, and fix RemoveImport shadowing

rewrite-python/rewrite/pyproject.toml

@@ -24,7 +24,7 @@ requires-python = ">=3.10"
 dependencies = [
     "cbor2>=5.6.5",
     "more_itertools>=10.0.0",
-    "ty-types>=0.0.19.dev20260223102528",  # Type inference CLI for Python type attribution
+    "ty-types>=0.0.19.dev20260223122104",  # Type inference CLI for Python type attribution
     "parso>=0.7.1,<0.8",  # Python 2/3 parser with CST support (0.8+ dropped Python 2.7 grammar)
 ]
 

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

@@ -20,7 +20,7 @@
 
 from rewrite.java import J
 from rewrite.java.support_types import JContainer, JRightPadded
-from rewrite.java.tree import Empty, FieldAccess, Identifier, Import, Space
+from rewrite.java.tree import Empty, FieldAccess, Identifier, Import, MethodDeclaration, Space
 from rewrite.markers import Markers
 from rewrite.python.tree import CompilationUnit, MultiImport
 from rewrite.python.visitor import PythonVisitor
@@ -129,6 +129,11 @@ def visit_multi_import(self, multi: MultiImport, p) -> J:
 
             def visit_identifier(self, ident: Identifier, p) -> J:
                 if not self.in_import:
+                    # Inside a function scope, identifiers with field_type are
+                    # local variables (shadowing the import), not actual uses.
+                    if self.cursor.first_enclosing(MethodDeclaration) is not None:
+                        if ident.field_type is not None:
+                            return ident
                     used.add(ident.simple_name)
                 return ident
 

rewrite-python/rewrite/src/rewrite/python/type_mapping.py

@@ -446,7 +446,12 @@ def _descriptor_to_java_type(self, descriptor: Dict[str, Any]) -> Optional[JavaT
 
         elif kind == 'classLiteral':
             class_name = descriptor.get('className', '')
-            class_type = self._create_class_type(class_name)
+            module_name = descriptor.get('moduleName')
+            if module_name and module_name != 'builtins':
+                fqn = f"{module_name}.{class_name}"
+            else:
+                fqn = class_name
+            class_type = self._create_class_type(fqn)
 
             # Infer Kind from supertypes before resolving them
             supertypes = descriptor.get('supertypes', [])
@@ -881,6 +886,10 @@ def _get_declaring_type(self, node: ast.Call) -> Optional[JavaType.FullyQualifie
                     kind = descriptor.get('kind')
                     if kind == 'module':
                         return self._create_class_type(descriptor.get('moduleName', ''))
+                    elif kind in ('function', 'boundMethod'):
+                        module_name = descriptor.get('moduleName')
+                        if module_name and module_name != 'builtins':
+                            return self._create_class_type(module_name)
 
         return self._infer_declaring_type_from_ast(node)
 

rewrite-python/rewrite/src/rewrite/test/rewrite_test.py

@@ -227,7 +227,7 @@ def _parse(self, specs: List[SourceSpec]) -> List[Tuple[SourceSpec, SourceFile]]
                 continue
 
             # Determine source path
-            source_path = spec.path or Path(f"{uuid4().hex}.{spec.ext}")
+            source_path = spec.path or Path(f"_{uuid4().hex}.{spec.ext}")
 
             # Parse the source
             source = dedent(spec.before)

rewrite-python/rewrite/tests/python/all/tree/assign_test.py

@@ -1,7 +1,3 @@
-import shutil
-
-import pytest
-
 from rewrite.java.support_types import JavaType
 from rewrite.java.tree import Assignment, Identifier, MethodInvocation
 from rewrite.python.tree import CompilationUnit
@@ -10,11 +6,6 @@
 
 Parameterized = JavaType.Parameterized
 
-requires_ty_cli = pytest.mark.skipif(
-    shutil.which('ty-types') is None,
-    reason="ty-types CLI is not installed"
-)
-
 
 def test_assign():
     # language=python
@@ -112,7 +103,6 @@ def test_assign_op():
     RecipeSpec().rewrite_run(python("a @= 1"))
 
 
-@requires_ty_cli
 def test_assign_type_attribution():
     """Verify that type attribution is populated on assignment AST nodes."""
     errors = []
@@ -146,7 +136,6 @@ def visit_assignment(self, assignment, p):
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
 
 
-@requires_ty_cli
 def test_assign_method_call_type_attribution():
     """Verify type attribution on an assignment from a method call like str.split()."""
     errors = []

rewrite-python/rewrite/tests/python/all/tree/binary_test.py

@@ -1,18 +1,9 @@
-import shutil
-
-import pytest
-
 from rewrite.java.support_types import JavaType
 from rewrite.java.tree import Binary
 from rewrite.python.tree import CompilationUnit
 from rewrite.python.visitor import PythonVisitor
 from rewrite.test import RecipeSpec, python
 
-requires_ty_cli = pytest.mark.skipif(
-    shutil.which('ty-types') is None,
-    reason="ty-types CLI is not installed"
-)
-
 
 def test_bool_ops():
     # language=python
@@ -108,7 +99,6 @@ def test_multiline_tuple_comparison():
     ),)
 
 
-@requires_ty_cli
 def test_arithmetic_type_attribution():
     """Verify that 1 + 2 has type Int."""
     errors = []
@@ -138,7 +128,6 @@ def visit_binary(self, binary, p):
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
 
 
-@requires_ty_cli
 def test_comparison_type_attribution():
     """Verify that 1 < 2 has type Boolean."""
     errors = []
@@ -168,7 +157,6 @@ def visit_binary(self, binary, p):
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
 
 
-@requires_ty_cli
 def test_boolean_op_type_attribution():
     """Verify that True and False has type Boolean."""
     errors = []

rewrite-python/rewrite/tests/python/all/tree/class_test.py

@@ -1,18 +1,11 @@
-import shutil
-
-import pytest
+from pathlib import Path
 
 from rewrite.java.support_types import JavaType
-from rewrite.java.tree import Assignment, Identifier
+from rewrite.java.tree import Assignment, ClassDeclaration, Identifier
 from rewrite.python.tree import CompilationUnit
 from rewrite.python.visitor import PythonVisitor
 from rewrite.test import RecipeSpec, python
 
-requires_ty_cli = pytest.mark.skipif(
-    shutil.which('ty-types') is None,
-    reason="ty-types CLI is not installed"
-)
-
 
 def test_empty():
     # language=python
@@ -125,7 +118,6 @@ class C3(Generic[T], metaclass=type, *[str]):
     ))
 
 
-@requires_ty_cli
 def test_generic_class_type_params():
     """Verify type parameters on a generic class like class Box[T]."""
     errors = []
@@ -174,9 +166,43 @@ def __init__(self, value: T) -> None:
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
 
 
-@requires_ty_cli
+def test_class_literal_module_qualified_fqn():
+    """Verify that a class defined in a module gets a module-qualified FQN on its classLiteral type."""
+    errors = []
+
+    def check_types(source_file):
+        assert isinstance(source_file, CompilationUnit)
+
+        class TypeChecker(PythonVisitor):
+            def visit_class_declaration(self, class_decl, p):
+                if not isinstance(class_decl, ClassDeclaration):
+                    return class_decl
+                cd_type = class_decl.type
+                if cd_type is None:
+                    errors.append("ClassDeclaration.type is None for MyClass")
+                elif isinstance(cd_type, JavaType.Class):
+                    fqn = cd_type._fully_qualified_name
+                    # The FQN should contain a module prefix (not just bare 'MyClass')
+                    if '.' not in fqn:
+                        errors.append(f"ClassDeclaration.type fqn '{fqn}' has no module prefix, expected '<module>.MyClass'")
+                return class_decl
+
+        TypeChecker().visit(source_file, None)
+
+    # language=python
+    RecipeSpec(type_attribution=True).rewrite_run(python(
+        """\
+        class MyClass:
+            pass
+        """,
+        path=Path("a.py"),
+        after_recipe=check_types,
+    ))
+    assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
+
+
 def test_class_instance_type_attribution():
-    """Verify that x = Foo() assigns a type with fqn 'Foo'."""
+    """Verify that x = Foo() assigns a type with fqn ending in 'Foo'."""
     errors = []
 
     def check_types(source_file):
@@ -189,8 +215,10 @@ def visit_assignment(self, assignment, p):
                 if assignment.type is None:
                     errors.append("Assignment.type is None for Foo()")
                 elif isinstance(assignment.type, JavaType.Class):
-                    if assignment.type._fully_qualified_name != 'Foo':
-                        errors.append(f"Assignment.type fqn is '{assignment.type._fully_qualified_name}', expected 'Foo'")
+                    fqn = assignment.type._fully_qualified_name
+                    # Accept 'Foo' or '<module>.Foo' (module-qualified with newer ty-types)
+                    if not fqn.endswith('Foo'):
+                        errors.append(f"Assignment.type fqn is '{fqn}', expected to end with 'Foo'")
                 else:
                     # Accept any non-None type
                     pass

rewrite-python/rewrite/tests/python/all/tree/collection_literal_test.py

@@ -1,8 +1,5 @@
-import shutil
 from typing import cast
 
-import pytest
-
 from rewrite.java import MethodDeclaration, Return
 from rewrite.java.support_types import JavaType
 from rewrite.java.tree import Assignment
@@ -12,11 +9,6 @@
 
 Parameterized = JavaType.Parameterized
 
-requires_ty_cli = pytest.mark.skipif(
-    shutil.which('ty-types') is None,
-    reason="ty-types CLI is not installed"
-)
-
 
 def test_empty_tuple():
     # language=python
@@ -107,7 +99,6 @@ def test_list_of_tuples_with_double_parens():
 """))
 
 
-@requires_ty_cli
 def test_list_literal_type_attribution():
     """Verify that [1, 2, 3] has type list."""
     errors = []
@@ -139,7 +130,6 @@ def visit_assignment(self, assignment, p):
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
 
 
-@requires_ty_cli
 def test_dict_literal_type_attribution():
     """Verify that {"a": 1} has type dict."""
     errors = []

rewrite-python/rewrite/tests/python/all/tree/def_test.py

@@ -1,18 +1,9 @@
-import shutil
-
-import pytest
-
 from rewrite.java.support_types import JavaType
 from rewrite.java.tree import MethodDeclaration
 from rewrite.python.tree import CompilationUnit
 from rewrite.python.visitor import PythonVisitor
 from rewrite.test import RecipeSpec, python
 
-requires_ty_cli = pytest.mark.skipif(
-    shutil.which('ty-types') is None,
-    reason="ty-types CLI is not installed"
-)
-
 
 def test_async_def():
     # language=python
@@ -24,7 +15,6 @@ async def main():
     ))
 
 
-@requires_ty_cli
 def test_async_def_type_attribution():
     """Verify that async def main() -> int has a method_type with a return_type.
 

rewrite-python/rewrite/tests/python/all/tree/field_access_test.py

@@ -1,17 +1,8 @@
-import shutil
-
-import pytest
-
 from rewrite.java.tree import FieldAccess
 from rewrite.python.tree import CompilationUnit
 from rewrite.python.visitor import PythonVisitor
 from rewrite.test import RecipeSpec, python
 
-requires_ty_cli = pytest.mark.skipif(
-    shutil.which('ty-types') is None,
-    reason="ty-types CLI is not installed"
-)
-
 
 # noinspection PyUnresolvedReferences
 def test_attribute():
@@ -25,7 +16,6 @@ def test_nested_attribute():
     RecipeSpec().rewrite_run(python("a = foo.bar.baz"))
 
 
-@requires_ty_cli
 def test_field_access_type_attribution():
     """Verify that os.path has a non-None FieldAccess.type."""
     errors = []

rewrite-python/rewrite/tests/python/all/tree/for_test.py

@@ -1,17 +1,8 @@
-import shutil
-
-import pytest
-
 from rewrite.java.tree import ForEachLoop, Identifier
 from rewrite.python.tree import CompilationUnit
 from rewrite.python.visitor import PythonVisitor
 from rewrite.test import RecipeSpec, python
 
-requires_ty_cli = pytest.mark.skipif(
-    shutil.which('ty-types') is None,
-    reason="ty-types CLI is not installed"
-)
-
 
 def test_for():
     # language=python
@@ -75,7 +66,6 @@ def test_async():
     ))
 
 
-@requires_ty_cli
 def test_for_loop_variable_type_attribution():
     """Verify that the loop control variable in 'for x in [1, 2, 3]' has a type."""
     errors = []

rewrite-python/rewrite/tests/python/all/tree/import_test.py

@@ -1,19 +1,10 @@
-import shutil
-
-import pytest
-
 from rewrite.java import Import
 from rewrite.java.support_types import JavaType
 from rewrite.java.tree import MethodInvocation
 from rewrite.python import CompilationUnit
 from rewrite.python.visitor import PythonVisitor
 from rewrite.test import RecipeSpec, python
 
-requires_ty_cli = pytest.mark.skipif(
-    shutil.which('ty-types') is None,
-    reason="ty-types CLI is not installed"
-)
-
 
 def _assert_single_j_import(cu: CompilationUnit) -> None:
     assert len(cu.statements) == 1
@@ -123,7 +114,6 @@ def visit_method_invocation(self, method, p):
     TypeChecker().visit(source_file, None)
 
 
-@requires_ty_cli
 def test_qualified_import_type_attribution():
     """import os; os.getcwd() → declaring_type.fqn == 'os'."""
     errors = []
@@ -138,8 +128,6 @@ def test_qualified_import_type_attribution():
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
 
 
-@requires_ty_cli
-@pytest.mark.xfail(reason="from-import direct calls do not yet populate declaring_type")
 def test_from_import_type_attribution():
     """from os import getcwd; getcwd() → declaring_type.fqn == 'os'."""
     errors = []
@@ -154,7 +142,6 @@ def test_from_import_type_attribution():
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
 
 
-@requires_ty_cli
 def test_aliased_import_type_attribution():
     """import os as o; o.getcwd() → declaring_type.fqn == 'os'."""
     errors = []
@@ -169,7 +156,6 @@ def test_aliased_import_type_attribution():
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
 
 
-@requires_ty_cli
 def test_aliased_from_import_type_attribution():
     """from os import getcwd as gwd; gwd() → declaring_type.fqn == 'os'."""
     errors = []