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)