Python: Add parameter field_type attribution and call-site type arguments (#6795)

* Add type attribution tests and fix method declaration/type hint types

- Add type attribution tests to 13 existing parser test files covering
  method invocations, binary ops, type hints, collection literals, imports,
  field access, class instances, method declarations, async defs, for loops,
  unary ops, ternaries, and lambdas
- Add `method_declaration_type()` to type_mapping.py to build JavaType.Method
  for function declarations using ty-types descriptor data with annotation fallback
- Add type attribution to ParameterizedType nodes in type hint expressions
- Fix pre-existing assign_test.py bug (simple_name access, FQN startswith)
- Add typing.Text test to test_type_attribution.py
- Bump ty-types dependency to >=0.0.19.dev20260223093555

* Add parameter field_type attribution and call-site type arguments

- Add param_type_info() to PythonTypeMapping for function parameter
  Identifier nodes to get JavaType.Variable field_type
- Update __convert_name() and map_arg() to flow field_type through to
  J.Identifier and NamedVariable
- Use call signature returnTypeId for call-site-specific return types
  (e.g. int for identity(42) instead of generic T)
- Populate _declared_formal_type_names on method invocation types from
  function descriptor type parameters
- Bump ty-types to 0.0.19.dev20260223102528 for callSignature
  typeArguments/returnTypeId support

Author: knutwannheden

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.dev20260223093555",  # Type inference CLI for Python type attribution
+    "ty-types>=0.0.19.dev20260223102528",  # 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/_parser_visitor.py

@@ -338,7 +338,8 @@ def visit_arguments(self, node, with_close_paren: bool = True) -> List[JRightPad
     def map_arg(self, node, default=None, vararg=False, kwarg=False):
         prefix = self.__source_before('**') if kwarg else self.__whitespace()
         vararg_prefix = self.__source_before('*') if vararg else None
-        name = self.__convert_name(node.arg, self._type_mapping.type(node))
+        expr_type, field_type = self._type_mapping.param_type_info(node)
+        name = self.__convert_name(node.arg, expr_type, field_type)
         after_name = self.__source_before(':') if node.annotation else Space.EMPTY
         type_expression = self.__convert_type(node.annotation) if node.annotation else None
         initializer = self.__pad_left(self.__source_before('='), self.__convert(default)) if default else None
@@ -359,7 +360,7 @@ def map_arg(self, node, default=None, vararg=False, kwarg=False):
                 cast(j.Identifier, name),
                 [],
                 initializer,
-                self.__as_variable_type(self._type_mapping.type(node))
+                field_type
             ), after_name)],
         )
 
@@ -3039,12 +3040,13 @@ def __as_method_type(t: Optional[JavaType]) -> Optional[JavaType]:
             return t
         return None
 
-    def __convert_name(self, name: str, name_type: Optional[JavaType] = None) -> NameTree:
+    def __convert_name(self, name: str, name_type: Optional[JavaType] = None,
+                        field_type: Optional[JavaType.Variable] = None) -> NameTree:
         def ident_or_field(parts: List[str]) -> NameTree:
             if len(parts) == 1:
                 space, actual_name = self.__consume_identifier(parts[-1])
                 return j.Identifier(random_id(), space, Markers.EMPTY, [], actual_name,
-                                    name_type, None)
+                                    name_type, field_type)
             else:
                 return j.FieldAccess(
                     random_id(),
@@ -3055,7 +3057,7 @@ def ident_or_field(parts: List[str]) -> NameTree:
                         self.__source_before('.'),
                         (lambda s, n: j.Identifier(random_id(), s, Markers.EMPTY, [], n,
                                      name_type,
-                                     None))(*self.__consume_identifier(parts[-1])),
+                                     field_type))(*self.__consume_identifier(parts[-1])),
                     ),
                     name_type
                 )

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

@@ -477,6 +477,17 @@ def _descriptor_to_java_type(self, descriptor: Dict[str, Any]) -> Optional[JavaT
                     if interfaces:
                         class_type._interfaces = interfaces
 
+            # Populate type parameters from typeVar descriptors
+            type_params = descriptor.get('typeParameters', [])
+            if type_params and getattr(class_type, '_type_parameters', None) is None:
+                resolved_type_params = []
+                for tp_id in type_params:
+                    tp_type = self._resolve_type(tp_id)
+                    if tp_type is not None:
+                        resolved_type_params.append(tp_type)
+                if resolved_type_params:
+                    class_type._type_parameters = resolved_type_params
+
             # Populate methods from function/boundMethod members
             members = descriptor.get('members', [])
             if members and getattr(class_type, '_methods', None) is None:
@@ -533,6 +544,10 @@ def _descriptor_to_java_type(self, descriptor: Dict[str, Any]) -> Optional[JavaT
         elif kind == 'property':
             return _UNKNOWN
 
+        elif kind == 'typeVar':
+            name = descriptor.get('name', '')
+            return self._create_class_type(name) if name else _UNKNOWN
+
         else:
             return _UNKNOWN
 
@@ -601,6 +616,16 @@ def name_type_info(self, node: ast.Name) -> Tuple[Optional[JavaType], Optional[J
             return expr_type, JavaType.Variable(_name=node.id, _type=expr_type)
         return expr_type, None
 
+    def param_type_info(self, node: ast.arg) -> Tuple[Optional[JavaType], Optional[JavaType.Variable]]:
+        """Get expression type and variable type for a function parameter.
+
+        Returns (expression_type, variable_field_type).
+        """
+        expr_type = self.type(node)
+        if expr_type is None:
+            return None, None
+        return expr_type, JavaType.Variable(_name=node.arg, _type=expr_type)
+
     def attribute_type_info(self, node: ast.Attribute,
                             receiver_type: Optional[JavaType] = None
                             ) -> Tuple[Optional[JavaType], Optional[JavaType.Variable]]:
@@ -660,7 +685,13 @@ def method_declaration_type(self, node: ast.FunctionDef) -> Optional[JavaType.Me
         if node.returns is not None:
             return_type = self.type(node.returns)
 
-        if not param_names and return_type is None:
+        # Extract type parameter names from Python 3.12+ type_params
+        type_param_names: List[str] = []
+        for tp in getattr(node, 'type_params', []) or []:
+            if hasattr(tp, 'name'):
+                type_param_names.append(tp.name)
+
+        if not param_names and return_type is None and not type_param_names:
             return None
 
         return JavaType.Method(
@@ -670,6 +701,7 @@ def method_declaration_type(self, node: ast.FunctionDef) -> Optional[JavaType.Me
             _return_type=return_type,
             _parameter_names=param_names if param_names else None,
             _parameter_types=param_types if param_types else None,
+            _declared_formal_type_names=type_param_names if type_param_names else None,
         )
 
     def _method_from_function_descriptor(
@@ -690,13 +722,16 @@ def _method_from_function_descriptor(
         if ret_id is not None:
             return_type = self._resolve_type(ret_id)
 
+        type_param_names = self._extract_type_param_names(descriptor)
+
         return JavaType.Method(
             _flags_bit_map=0,
             _declaring_type=None,
             _name=name,
             _return_type=return_type,
             _parameter_names=param_names if param_names else None,
             _parameter_types=param_types if param_types else None,
+            _declared_formal_type_names=type_param_names if type_param_names else None,
         )
 
     def method_invocation_type(self, node: ast.Call) -> Optional[JavaType.Method]:
@@ -729,13 +764,22 @@ def method_invocation_type(self, node: ast.Call) -> Optional[JavaType.Method]:
         # Get return type
         return_type = self._get_return_type(node)
 
+        # Extract type parameter names from function descriptor
+        type_param_names: List[str] = []
+        func_type_id = self._lookup_func_type_id(node)
+        if func_type_id is not None:
+            func_desc = self._type_registry.get(func_type_id)
+            if func_desc:
+                type_param_names = self._extract_type_param_names(func_desc)
+
         return JavaType.Method(
             _flags_bit_map=0,
             _declaring_type=declaring_type,
             _name=method_name,
             _return_type=return_type,
             _parameter_names=param_names if param_names else None,
             _parameter_types=param_types if param_types else None,
+            _declared_formal_type_names=type_param_names if type_param_names else None,
         )
 
     def _extract_method_name(self, node: ast.Call) -> Optional[str]:
@@ -1037,9 +1081,20 @@ def _get_parameter_types(self, node: ast.Call) -> Optional[List[JavaType]]:
     def _get_return_type(self, node: ast.Call) -> Optional[JavaType]:
         """Get the return type of a method call.
 
-        First tries the ExprCall node type (which IS the return type),
-        then falls back to the function descriptor's returnType field.
+        Prefers the call-site-specific returnTypeId from the call signature
+        (which gives resolved types like int instead of generic T),
+        then tries the ExprCall node type, then falls back to the function
+        descriptor's returnType field.
         """
+        # Prefer call-site-specific return type from call signature
+        sig = self._lookup_call_signature(node)
+        if sig:
+            ret_id = sig.get('returnTypeId')
+            if ret_id is not None:
+                result = self._resolve_type(ret_id)
+                if result is not None:
+                    return result
+
         # The type of an ExprCall node in ty-types IS the return type
         type_id = self._lookup_type_id(node)
         if type_id is not None:
@@ -1083,15 +1138,29 @@ def _create_method_from_descriptor(self, descriptor: Dict[str, Any],
             else:
                 param_types.append(_UNKNOWN)
 
+        type_param_names = self._extract_type_param_names(descriptor)
+
         return JavaType.Method(
             _flags_bit_map=0,
             _declaring_type=declaring_type,
             _name=name,
             _return_type=return_type,
             _parameter_names=param_names if param_names else None,
             _parameter_types=param_types if param_types else None,
+            _declared_formal_type_names=type_param_names if type_param_names else None,
         )
 
+    def _extract_type_param_names(self, descriptor: Dict[str, Any]) -> List[str]:
+        """Extract type parameter names from a descriptor's typeParameters list."""
+        names: List[str] = []
+        for tp_id in descriptor.get('typeParameters', []):
+            tp_desc = self._type_registry.get(tp_id)
+            if tp_desc and tp_desc.get('kind') == 'typeVar':
+                name = tp_desc.get('name', '')
+                if name:
+                    names.append(name)
+        return names
+
     def _create_class_type(self, fqn: str) -> JavaType.Class:
         """Create a JavaType.Class from a fully qualified name."""
         if fqn in self._type_cache:

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

@@ -3,7 +3,7 @@
 import pytest
 
 from rewrite.java.support_types import JavaType
-from rewrite.java.tree import Assignment
+from rewrite.java.tree import Assignment, Identifier
 from rewrite.python.tree import CompilationUnit
 from rewrite.python.visitor import PythonVisitor
 from rewrite.test import RecipeSpec, python
@@ -125,6 +125,55 @@ 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 = []
+
+    def check_types(source_file):
+        assert isinstance(source_file, CompilationUnit)
+
+        class TypeChecker(PythonVisitor):
+            def visit_assignment(self, assignment, p):
+                if not isinstance(assignment, Assignment):
+                    return assignment
+                # Only check the `x = Box(42)` assignment
+                if not isinstance(assignment.variable, Identifier) or assignment.variable.simple_name != 'x':
+                    return assignment
+                if assignment.type is None:
+                    errors.append("Assignment.type is None for x = Box(42)")
+                elif isinstance(assignment.type, JavaType.Class):
+                    if assignment.type._fully_qualified_name != 'Box':
+                        errors.append(f"Assignment.type fqn is '{assignment.type._fully_qualified_name}', expected 'Box'")
+                    type_params = getattr(assignment.type, '_type_parameters', None)
+                    if type_params is None:
+                        errors.append("Box class type has no _type_parameters")
+                    elif len(type_params) != 1:
+                        errors.append(f"Box class type has {len(type_params)} type params, expected 1")
+                    else:
+                        tp = type_params[0]
+                        if isinstance(tp, JavaType.Class):
+                            if tp._fully_qualified_name != 'T':
+                                errors.append(f"type_parameter fqn is '{tp._fully_qualified_name}', expected 'T'")
+                        else:
+                            errors.append(f"type_parameter is {type(tp).__name__}, expected Class")
+                return assignment
+
+        TypeChecker().visit(source_file, None)
+
+    # language=python
+    RecipeSpec(type_attribution=True).rewrite_run(python(
+        """\
+        class Box[T]:
+            def __init__(self, value: T) -> None:
+                self.value = value
+        x = Box(42)
+        """,
+        after_recipe=check_types,
+    ))
+    assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
+
+
 @requires_ty_cli
 def test_class_instance_type_attribution():
     """Verify that x = Foo() assigns a type with fqn 'Foo'."""

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

@@ -3,7 +3,7 @@
 import pytest
 
 from rewrite.java.support_types import JavaType
-from rewrite.java.tree import MethodDeclaration
+from rewrite.java.tree import MethodDeclaration, VariableDeclarations
 from rewrite.python.tree import CompilationUnit
 from rewrite.python.visitor import PythonVisitor
 from rewrite.test import RecipeSpec, python
@@ -167,6 +167,45 @@ def f(
     ))
 
 
+@requires_ty_cli
+def test_generic_function_type_params():
+    """Verify method_type.declared_formal_type_names for def identity[T](x: T) -> T."""
+    errors = []
+
+    def check_types(source_file):
+        assert isinstance(source_file, CompilationUnit)
+
+        class TypeChecker(PythonVisitor):
+            def visit_method_declaration(self, method, p):
+                if not isinstance(method, MethodDeclaration):
+                    return method
+                if method.name.simple_name != 'identity':
+                    return method
+                if method.method_type is None:
+                    errors.append("MethodDeclaration.method_type is None")
+                else:
+                    mt = method.method_type
+                    if mt._declared_formal_type_names is None:
+                        errors.append("method_type._declared_formal_type_names is None")
+                    elif mt._declared_formal_type_names != ['T']:
+                        errors.append(f"method_type._declared_formal_type_names is {mt._declared_formal_type_names}, expected ['T']")
+                    if mt._parameter_names is not None and 'x' not in mt._parameter_names:
+                        errors.append(f"parameter_names {mt._parameter_names} does not contain 'x'")
+                return method
+
+        TypeChecker().visit(source_file, None)
+
+    # language=python
+    RecipeSpec(type_attribution=True).rewrite_run(python(
+        """\
+        def identity[T](x: T) -> T:
+            return x
+        """,
+        after_recipe=check_types,
+    ))
+    assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
+
+
 @requires_ty_cli
 def test_method_declaration_type_attribution():
     """Verify method_type on a function with typed parameters and return type."""
@@ -216,3 +255,43 @@ def foo(a: int, b: str) -> bool:
         after_recipe=check_types,
     ))
     assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)
+
+
+@requires_ty_cli
+def test_param_identifier_field_type():
+    """Verify J.Identifier.field_type is JavaType.Variable for typed parameters."""
+    errors = []
+
+    def check_types(source_file):
+        assert isinstance(source_file, CompilationUnit)
+
+        class TypeChecker(PythonVisitor):
+            def visit_variable_declarations(self, vd, p):
+                if not isinstance(vd, VariableDeclarations):
+                    return vd
+                for named_var in vd.variables:
+                    ident = named_var.name
+                    if ident.simple_name not in ('x', 'y'):
+                        continue
+                    if ident.field_type is None:
+                        errors.append(f"Identifier '{ident.simple_name}' has field_type=None")
+                    elif not isinstance(ident.field_type, JavaType.Variable):
+                        errors.append(f"Identifier '{ident.simple_name}' field_type is {type(ident.field_type)}, expected JavaType.Variable")
+                    else:
+                        if ident.field_type._name != ident.simple_name:
+                            errors.append(f"field_type._name is '{ident.field_type._name}', expected '{ident.simple_name}'")
+                        if ident.field_type._type is None:
+                            errors.append(f"field_type._type is None for '{ident.simple_name}'")
+                return vd
+
+        TypeChecker().visit(source_file, None)
+
+    # language=python
+    RecipeSpec(type_attribution=True).rewrite_run(python(
+        """\
+        def greet(x: int, y: str) -> bool:
+            return True
+        """,
+        after_recipe=check_types,
+    ))
+    assert not errors, "Type attribution errors:\n" + "\n".join(f"  - {e}" for e in errors)