Python: preserve PEP 695 type-alias type parameters across RPC

rewrite-python/rewrite/src/rewrite/rpc/python_receiver.py

@@ -323,9 +323,10 @@ def _visit_comprehension_clause(self, cc: ComprehensionExpression.Clause, q: Rpc
 
     def _visit_type_alias(self, ta: TypeAlias, q: RpcReceiveQueue) -> TypeAlias:
         name = q.receive(ta.name)
+        type_parameters = q.receive(ta.padding.type_parameters, lambda c: self._receive_container(c, q) if c else None)
         value = q.receive(ta.padding.value)
         type_ = q.receive(ta.type)
-        return replace_if_changed(ta, name=name, value=value, type=type_)
+        return replace_if_changed(ta, name=name, type_parameters=type_parameters, value=value, type=type_)
 
     def _visit_yield_from(self, yf: YieldFrom, q: RpcReceiveQueue) -> YieldFrom:
         expression = q.receive(yf.expression)

rewrite-python/rewrite/src/rewrite/rpc/python_sender.py

@@ -290,6 +290,7 @@ def _visit_comprehension_clause(self, cc: ComprehensionExpression.Clause, q: 'Rp
 
     def _visit_type_alias(self, ta: TypeAlias, q: 'RpcSendQueue') -> None:
         q.get_and_send(ta, lambda x: x.name, lambda el: self._visit(el, q))
+        q.get_and_send(ta, lambda x: x.padding.type_parameters, lambda c: self._visit_container(c, q) if c else None)
         q.get_and_send(ta, lambda x: x.padding.value, lambda el: self._visit_left_padded(el, q))
         q.get_and_send_as_ref(ta, lambda x: x.type, lambda t: self._visit_type(t, q) if t else None)
 

rewrite-python/rewrite/tests/rpc/test_type_alias_roundtrip.py

@@ -0,0 +1,58 @@
+# Copyright 2025 the original author or authors.
+#
+# Licensed under the Moderne Source Available License (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://docs.moderne.io/licensing/moderne-source-available-license
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PEP 695 type-alias type parameters must survive the RPC round trip.
+
+All four codecs (Python + Java, sender + receiver) previously omitted
+TypeAlias.type_parameters, silently narrowing `type X[T] = ...` to `type X = ...`.
+The loss is invisible to the in-process printer suite and baked into every LST,
+so only a serialize/deserialize round trip catches it.
+"""
+import ast
+
+import pytest
+
+from rewrite.python._parser_visitor import ParserVisitor
+from rewrite.python.printer import PythonPrinter
+from rewrite.rpc.python_receiver import PythonRpcReceiver
+from rewrite.rpc.receive_queue import RpcReceiveQueue
+from rewrite.rpc.send_queue import RpcSendQueue
+
+_CU_TYPE = "org.openrewrite.python.tree.Py$CompilationUnit"
+
+
+def _round_trip_print(src):
+    cu = ParserVisitor(src, "<m>", None).visit_Module(ast.parse(src))
+    data = list(RpcSendQueue(_CU_TYPE).generate(cu, None))
+
+    def pull():
+        out = data[:]
+        data.clear()
+        return out
+
+    rebuilt = PythonRpcReceiver().receive(None, RpcReceiveQueue({}, _CU_TYPE, pull))
+    return PythonPrinter().print(cu), PythonPrinter().print(rebuilt)
+
+
+# Spaced out to also exercise whitespace fidelity: JContainer.before (before `[`) and the
+# element padding (inside the brackets, around the `:` bound and the commas).
+@pytest.mark.parametrize("src", [
+    "type X [ T ] = list[T]\n",
+    "type X [ T : int ] = list[T]\n",
+    "type F [ T , *Ts , **P ] = Callable[P, T]\n",
+    "type X = int\n",  # no params — must remain a no-op
+])
+def test_type_alias_params_survive_rpc(src):
+    before, after = _round_trip_print(src)
+    assert after == before

rewrite-python/src/main/java/org/openrewrite/python/internal/rpc/PythonReceiver.java

@@ -237,6 +237,7 @@ public J visitComprehensionClause(Py.ComprehensionExpression.Clause clause, RpcR
     public J visitTypeAlias(Py.TypeAlias typeAlias, RpcReceiveQueue q) {
         return typeAlias
                 .withName(q.receive(typeAlias.getName(), name -> (J.Identifier) visitNonNull(name, q)))
+                .getPadding().withTypeParameters(q.receive(typeAlias.getPadding().getTypeParameters(), el -> visitContainer(el, q)))
                 .getPadding().withValue(q.receive(typeAlias.getPadding().getValue(), el -> visitLeftPadded(el, q)))
                 .withType(q.receive(typeAlias.getType(), type -> visitType(type, q)));
     }

rewrite-python/src/main/java/org/openrewrite/python/internal/rpc/PythonSender.java

@@ -231,6 +231,7 @@ public J visitComprehensionClause(Py.ComprehensionExpression.Clause clause, RpcS
     @Override
     public J visitTypeAlias(Py.TypeAlias typeAlias, RpcSendQueue q) {
         q.getAndSend(typeAlias, Py.TypeAlias::getName, el -> visit(el, q));
+        q.getAndSend(typeAlias, el -> el.getPadding().getTypeParameters(), el -> visitContainer(el, q));
         q.getAndSend(typeAlias, el -> el.getPadding().getValue(), el -> visitLeftPadded(el, q));
         q.getAndSend(typeAlias, el -> asRef(el.getType()), el -> visitType(getValueNonNull(el), q));
         return typeAlias;

rewrite-python/src/main/java/org/openrewrite/python/tree/Py.java

@@ -1562,6 +1562,17 @@ final class TypeAlias implements Py, Statement, TypedTree {
         @Getter
         J.Identifier name;
 
+        @Nullable
+        JContainer<J.TypeParameter> typeParameters;
+
+        public @Nullable List<J.TypeParameter> getTypeParameters() {
+            return typeParameters == null ? null : typeParameters.getElements();
+        }
+
+        public TypeAlias withTypeParameters(@Nullable List<J.TypeParameter> typeParameters) {
+            return getPadding().withTypeParameters(JContainer.withElementsNullable(this.typeParameters, typeParameters));
+        }
+
         JLeftPadded<J> value;
 
         public J getValue() {
@@ -1612,12 +1623,20 @@ public String toString() {
         public static class Padding {
             private final TypeAlias t;
 
+            public @Nullable JContainer<J.TypeParameter> getTypeParameters() {
+                return t.typeParameters;
+            }
+
+            public TypeAlias withTypeParameters(@Nullable JContainer<J.TypeParameter> typeParameters) {
+                return t.typeParameters == typeParameters ? t : new TypeAlias(t.id, t.prefix, t.markers, t.name, typeParameters, t.value, t.type);
+            }
+
             public JLeftPadded<J> getValue() {
                 return t.value;
             }
 
             public TypeAlias withValue(JLeftPadded<J> assignment) {
-                return t.value == assignment ? t : new TypeAlias(t.id, t.prefix, t.markers, t.name, assignment, t.type);
+                return t.value == assignment ? t : new TypeAlias(t.id, t.prefix, t.markers, t.name, t.typeParameters, assignment, t.type);
             }
         }
     }

rewrite-python/src/test/java/org/openrewrite/python/internal/rpc/PythonSenderReceiverRoundTripTest.java

@@ -19,6 +19,9 @@
 import org.junit.jupiter.api.Test;
 import org.openrewrite.Tree;
 import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JContainer;
+import org.openrewrite.java.tree.JLeftPadded;
+import org.openrewrite.java.tree.JRightPadded;
 import org.openrewrite.java.tree.JavaType;
 import org.openrewrite.java.tree.Space;
 import org.openrewrite.marker.Markers;
@@ -106,6 +109,37 @@ void expressionStatementWithAwaitAddRoundTrip() {
         assertThat(received.getExpression()).isInstanceOf(Py.Await.class);
     }
 
+    @Test
+    void typeAliasTypeParametersRoundTrip() {
+        // PEP 695 `type X[T] = int` — all four codecs previously dropped TypeAlias.typeParameters.
+        J.TypeParameter typeParam = new J.TypeParameter(
+                Tree.randomId(), Space.EMPTY, Markers.EMPTY,
+                Collections.emptyList(), Collections.emptyList(), ident("T"), null);
+        Py.TypeAlias after = new Py.TypeAlias(
+                Tree.randomId(), Space.EMPTY, Markers.EMPTY, ident("X"),
+                JContainer.build(Space.EMPTY, List.of(JRightPadded.build(typeParam)), Markers.EMPTY),
+                JLeftPadded.build((J) ident("int")).withBefore(Space.SINGLE_SPACE),
+                null);
+
+        sq.send(after, null, null);
+        sq.flush();
+
+        Py.TypeAlias received = rq.receive(null);
+
+        assertThat(received).isNotNull();
+        assertThat(received.getTypeParameters())
+                .as("PEP 695 type parameters must survive the RPC round trip")
+                .isNotNull()
+                .hasSize(1);
+        assertThat(((J.Identifier) received.getTypeParameters().get(0).getName()).getSimpleName())
+                .isEqualTo("T");
+    }
+
+    private static J.Identifier ident(String name) {
+        return new J.Identifier(Tree.randomId(), Space.EMPTY, Markers.EMPTY,
+                Collections.emptyList(), name, null, null);
+    }
+
     private static Py.ExpressionStatement expressionStatementOfAwait(String name, Space awaitPrefix) {
         UUID esId = Tree.randomId();
         UUID awaitId = Tree.randomId();