Add Py.ChainedAssignment

Not yet used by the parser, but that part will be implemented soon.
openrewrite · Oct 17, 2024 · 12fe8a3 · 12fe8a3
1 parent e71b20a
commit 12fe8a3
Show file tree

Hide file tree

Showing 13 changed files with 284 additions and 12 deletions.
diff --git a/rewrite-python-remote/src/main/java/org/openrewrite/python/remote/PythonReceiver.java b/rewrite-python-remote/src/main/java/org/openrewrite/python/remote/PythonReceiver.java
@@ -93,6 +93,17 @@ public Py.Binary visitBinary(Py.Binary binary, ReceiverContext ctx) {
             return binary;
         }
 
+        @Override
+        public Py.ChainedAssignment visitChainedAssignment(Py.ChainedAssignment chainedAssignment, ReceiverContext ctx) {
+            chainedAssignment = chainedAssignment.withId(ctx.receiveNonNullValue(chainedAssignment.getId(), UUID.class));
+            chainedAssignment = chainedAssignment.withPrefix(ctx.receiveNonNullNode(chainedAssignment.getPrefix(), PythonReceiver::receiveSpace));
+            chainedAssignment = chainedAssignment.withMarkers(ctx.receiveNonNullNode(chainedAssignment.getMarkers(), ctx::receiveMarkers));
+            chainedAssignment = chainedAssignment.getPadding().withVariables(ctx.receiveNonNullNodes(chainedAssignment.getPadding().getVariables(), PythonReceiver::receiveRightPaddedTree));
+            chainedAssignment = chainedAssignment.withAssignment(ctx.receiveNonNullNode(chainedAssignment.getAssignment(), ctx::receiveTree));
+            chainedAssignment = chainedAssignment.withType(ctx.receiveValue(chainedAssignment.getType(), JavaType.class));
+            return chainedAssignment;
+        }
+
         @Override
         public Py.ExceptionType visitExceptionType(Py.ExceptionType exceptionType, ReceiverContext ctx) {
             exceptionType = exceptionType.withId(ctx.receiveNonNullValue(exceptionType.getId(), UUID.class));
@@ -1123,6 +1134,17 @@ public <T> T create(Class<T> type, ReceiverContext ctx) {
                 );
             }
 
+            if (type == Py.ChainedAssignment.class) {
+                return (T) new Py.ChainedAssignment(
+                    ctx.receiveNonNullValue(null, UUID.class),
+                    ctx.receiveNonNullNode(null, PythonReceiver::receiveSpace),
+                    ctx.receiveNonNullNode(null, ctx::receiveMarkers),
+                    ctx.receiveNonNullNodes(null, PythonReceiver::receiveRightPaddedTree),
+                    ctx.receiveNonNullNode(null, ctx::receiveTree),
+                    ctx.receiveValue(null, JavaType.class)
+                );
+            }
+
             if (type == Py.ExceptionType.class) {
                 return (T) new Py.ExceptionType(
                     ctx.receiveNonNullValue(null, UUID.class),

diff --git a/rewrite-python-remote/src/main/java/org/openrewrite/python/remote/PythonSender.java b/rewrite-python-remote/src/main/java/org/openrewrite/python/remote/PythonSender.java
@@ -79,6 +79,17 @@ public Py.Binary visitBinary(Py.Binary binary, SenderContext ctx) {
             return binary;
         }
 
+        @Override
+        public Py.ChainedAssignment visitChainedAssignment(Py.ChainedAssignment chainedAssignment, SenderContext ctx) {
+            ctx.sendValue(chainedAssignment, Py.ChainedAssignment::getId);
+            ctx.sendNode(chainedAssignment, Py.ChainedAssignment::getPrefix, PythonSender::sendSpace);
+            ctx.sendNode(chainedAssignment, Py.ChainedAssignment::getMarkers, ctx::sendMarkers);
+            ctx.sendNodes(chainedAssignment, e -> e.getPadding().getVariables(), PythonSender::sendRightPadded, e -> e.getElement().getId());
+            ctx.sendNode(chainedAssignment, Py.ChainedAssignment::getAssignment, ctx::sendTree);
+            ctx.sendTypedValue(chainedAssignment, Py.ChainedAssignment::getType);
+            return chainedAssignment;
+        }
+
         @Override
         public Py.ExceptionType visitExceptionType(Py.ExceptionType exceptionType, SenderContext ctx) {
             ctx.sendValue(exceptionType, Py.ExceptionType::getId);

diff --git a/rewrite-python/src/main/java/org/openrewrite/python/PythonIsoVisitor.java b/rewrite-python/src/main/java/org/openrewrite/python/PythonIsoVisitor.java
@@ -27,6 +27,21 @@ public Py.Await visitAwait(Py.Await ogAwait, P p) {
         return (Py.Await) super.visitAwait(ogAwait, p);
     }
 
+    @Override
+    public Py.Binary visitBinary(Py.Binary binary, P p) {
+        return (Py.Binary) super.visitBinary(binary, p);
+    }
+
+    @Override
+    public Py.ChainedAssignment visitChainedAssignment(Py.ChainedAssignment chainedAssignment, P p) {
+        return (Py.ChainedAssignment) super.visitChainedAssignment(chainedAssignment, p);
+    }
+
+    @Override
+    public Py.CollectionLiteral visitCollectionLiteral(Py.CollectionLiteral coll, P p) {
+        return (Py.CollectionLiteral) super.visitCollectionLiteral(coll, p);
+    }
+
     @Override
     public Py.CompilationUnit visitCompilationUnit(Py.CompilationUnit cu, P p) {
         return (Py.CompilationUnit) super.visitCompilationUnit(cu, p);
@@ -52,11 +67,6 @@ public Py.Del visitDel(Py.Del ogDel, P p) {
         return (Py.Del) super.visitDel(ogDel, p);
     }
 
-    @Override
-    public Py.CollectionLiteral visitCollectionLiteral(Py.CollectionLiteral coll, P p) {
-        return (Py.CollectionLiteral) super.visitCollectionLiteral(coll, p);
-    }
-
     @Override
     public Py.DictLiteral visitDictLiteral(Py.DictLiteral dict, P p) {
         return (Py.DictLiteral) super.visitDictLiteral(dict, p);
@@ -227,11 +237,6 @@ public J.Binary visitBinary(J.Binary binary, P p) {
         return (J.Binary) super.visitBinary(binary, p);
     }
 
-    @Override
-    public Py.Binary visitBinary(Py.Binary binary, P p) {
-        return (Py.Binary) super.visitBinary(binary, p);
-    }
-
     @Override
     public J.Block visitBlock(J.Block block, P p) {
         return (J.Block) super.visitBlock(block, p);

diff --git a/rewrite-python/src/main/java/org/openrewrite/python/PythonVisitor.java b/rewrite-python/src/main/java/org/openrewrite/python/PythonVisitor.java
@@ -66,6 +66,21 @@ public J visitBinary(Py.Binary binary, P p) {
         return b;
     }
 
+    public J visitChainedAssignment(Py.ChainedAssignment chainedAssignment, P p) {
+        Py.ChainedAssignment ca = chainedAssignment;
+        ca = ca.withPrefix(visitSpace(ca.getPrefix(), PySpace.Location.CHAINED_ASSIGNMENT_PREFIX, p));
+        ca = ca.withMarkers(visitMarkers(ca.getMarkers(), p));
+        Statement temp = (Statement) visitStatement(ca, p);
+        if (!(temp instanceof Py.ChainedAssignment)) {
+            return temp;
+        } else {
+            ca = (Py.ChainedAssignment) temp;
+        }
+        ca = ca.getPadding().withVariables(ListUtils.map(ca.getPadding().getVariables(), t -> visitRightPadded(t, PyRightPadded.Location.CHAINED_ASSIGNMENT_VARIABLE, p)));
+        ca = ca.withAssignment(visitAndCast(ca.getAssignment(), p));
+        return ca;
+    }
+
     public J visitCollectionLiteral(Py.CollectionLiteral coll, P p) {
         Py.CollectionLiteral c = coll;
         c = c.withPrefix(visitSpace(c.getPrefix(), PySpace.Location.COLLECTION_LITERAL_PREFIX, p));

diff --git a/rewrite-python/src/main/java/org/openrewrite/python/internal/PythonPrinter.java b/rewrite-python/src/main/java/org/openrewrite/python/internal/PythonPrinter.java
@@ -136,6 +136,15 @@ public J visitBinary(Py.Binary binary, PrintOutputCapture<P> p) {
         return binary;
     }
 
+    @Override
+    public J visitChainedAssignment(Py.ChainedAssignment chainedAssignment, PrintOutputCapture<P> p) {
+        beforeSyntax(chainedAssignment, PySpace.Location.CHAINED_ASSIGNMENT_PREFIX, p);
+        visitRightPadded(chainedAssignment.getPadding().getVariables(), PyRightPadded.Location.CHAINED_ASSIGNMENT_VARIABLE, "=", p);
+        visit(chainedAssignment.getAssignment(), p);
+        afterSyntax(chainedAssignment, p);
+        return chainedAssignment;
+    }
+
     @Override
     public J visitCollectionLiteral(Py.CollectionLiteral coll, PrintOutputCapture<P> p) {
         beforeSyntax(coll, PySpace.Location.COLLECTION_LITERAL_PREFIX, p);

diff --git a/rewrite-python/src/main/java/org/openrewrite/python/tree/Py.java b/rewrite-python/src/main/java/org/openrewrite/python/tree/Py.java
@@ -201,6 +201,84 @@ public Py.Binary withOperator(JLeftPadded<Type> operator) {
         }
     }
 
+    @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
+    @EqualsAndHashCode(callSuper = false, onlyExplicitlyIncluded = true)
+    @RequiredArgsConstructor
+    @AllArgsConstructor(access = AccessLevel.PRIVATE)
+    @Data
+    final class ChainedAssignment implements Py, Statement, TypedTree {
+
+        @Nullable
+        @NonFinal
+        transient WeakReference<Padding> padding;
+
+        @With
+        @EqualsAndHashCode.Include
+        UUID id;
+
+        @With
+        Space prefix;
+
+        @With
+        Markers markers;
+
+        List<JRightPadded<Expression>> variables;
+
+        public List<Expression> getVariables() {
+            return JRightPadded.getElements(variables);
+        }
+
+        public ChainedAssignment withVariables(List<Expression> variables) {
+            return getPadding().withVariables(JRightPadded.withElements(this.variables, variables));
+        }
+
+        @With
+        Expression assignment;
+
+        @With
+        @Nullable
+        JavaType type;
+
+        @Override
+        public <P> J acceptPython(PythonVisitor<P> v, P p) {
+            return v.visitChainedAssignment(this, p);
+        }
+
+        @Transient
+        @Override
+        public CoordinateBuilder.Statement getCoordinates() {
+            return new CoordinateBuilder.Statement(this);
+        }
+
+        public Padding getPadding() {
+            Padding p;
+            if (this.padding == null) {
+                p = new Padding(this);
+                this.padding = new WeakReference<>(p);
+            } else {
+                p = this.padding.get();
+                if (p == null || p.t != this) {
+                    p = new Padding(this);
+                    this.padding = new WeakReference<>(p);
+                }
+            }
+            return p;
+        }
+
+        @RequiredArgsConstructor
+        public static class Padding {
+            private final ChainedAssignment t;
+
+            public List<JRightPadded<Expression>> getVariables() {
+                return t.variables;
+            }
+
+            public ChainedAssignment withVariables(List<JRightPadded<Expression>> variables) {
+                return t.variables == variables ? t : new ChainedAssignment(t.id, t.prefix, t.markers, variables, t.assignment, t.type);
+            }
+        }
+    }
+
     @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
     @EqualsAndHashCode(callSuper = false)
     @Data

diff --git a/rewrite-python/src/main/java/org/openrewrite/python/tree/PyRightPadded.java b/rewrite-python/src/main/java/org/openrewrite/python/tree/PyRightPadded.java
@@ -23,6 +23,7 @@ public class PyRightPadded {
     @RequiredArgsConstructor
     public enum Location {
         ASSERT_ELEMENT(PySpace.Location.ASSERT_ELEMENT_SUFFIX),
+        CHAINED_ASSIGNMENT_VARIABLE(PySpace.Location.CHAINED_ASSIGNMENT_VARIABLE_SUFFIX),
         COLLECTION_LITERAL_ELEMENT(PySpace.Location.COLLECTION_LITERAL_ELEMENT_SUFFIX),
         DEL_ELEMENT(PySpace.Location.DEL_ELEMENT_SUFFIX),
         DICT_ENTRY_KEY(PySpace.Location.DICT_ENTRY_KEY_SUFFIX),

diff --git a/rewrite-python/src/main/java/org/openrewrite/python/tree/PySpace.java b/rewrite-python/src/main/java/org/openrewrite/python/tree/PySpace.java
@@ -322,6 +322,8 @@ public enum Location {
         BINARY_NEGATION,
         BINARY_OPERATOR,
         BINARY_PREFIX,
+        CHAINED_ASSIGNMENT_PREFIX,
+        CHAINED_ASSIGNMENT_VARIABLE_SUFFIX,
         COLLECTION_LITERAL_ELEMENT_SUFFIX,
         COLLECTION_LITERAL_PREFIX,
         COMPREHENSION_CLAUSE_PREFIX,

diff --git a/rewrite/rewrite/python/remote/receiver.py b/rewrite/rewrite/python/remote/receiver.py
@@ -51,6 +51,15 @@ def visit_python_binary(self, binary: Binary, ctx: ReceiverContext) -> J:
             binary = binary.with_type(ctx.receive_value(binary.type, JavaType))
             return binary
 
+        def visit_chained_assignment(self, chained_assignment: ChainedAssignment, ctx: ReceiverContext) -> J:
+            chained_assignment = chained_assignment.with_id(ctx.receive_value(chained_assignment.id, UUID))
+            chained_assignment = chained_assignment.with_prefix(ctx.receive_node(chained_assignment.prefix, PythonReceiver.receive_space))
+            chained_assignment = chained_assignment.with_markers(ctx.receive_node(chained_assignment.markers, ctx.receive_markers))
+            chained_assignment = chained_assignment.padding.with_variables(ctx.receive_nodes(chained_assignment.padding.variables, PythonReceiver.receive_right_padded_tree))
+            chained_assignment = chained_assignment.with_assignment(ctx.receive_node(chained_assignment.assignment, ctx.receive_tree))
+            chained_assignment = chained_assignment.with_type(ctx.receive_value(chained_assignment.type, JavaType))
+            return chained_assignment
+
         def visit_exception_type(self, exception_type: ExceptionType, ctx: ReceiverContext) -> J:
             exception_type = exception_type.with_id(ctx.receive_value(exception_type.id, UUID))
             exception_type = exception_type.with_prefix(ctx.receive_node(exception_type.prefix, PythonReceiver.receive_space))
@@ -919,6 +928,16 @@ def create(self, type: str, ctx: ReceiverContext) -> Tree:
                     ctx.receive_value(None, JavaType)
                 )
 
+            if type in ["rewrite.python.tree.ChainedAssignment", "org.openrewrite.python.tree.Py$ChainedAssignment"]:
+                return ChainedAssignment(
+                    ctx.receive_value(None, UUID),
+                    ctx.receive_node(None, PythonReceiver.receive_space),
+                    ctx.receive_node(None, ctx.receive_markers),
+                    ctx.receive_nodes(None, PythonReceiver.receive_right_padded_tree),
+                    ctx.receive_node(None, ctx.receive_tree),
+                    ctx.receive_value(None, JavaType)
+                )
+
             if type in ["rewrite.python.tree.ExceptionType", "org.openrewrite.python.tree.Py$ExceptionType"]:
                 return ExceptionType(
                     ctx.receive_value(None, UUID),

diff --git a/rewrite/rewrite/python/remote/sender.py b/rewrite/rewrite/python/remote/sender.py
@@ -45,6 +45,15 @@ def visit_python_binary(self, binary: Binary, ctx: SenderContext) -> J:
             ctx.send_typed_value(binary, attrgetter('_type'))
             return binary
 
+        def visit_chained_assignment(self, chained_assignment: ChainedAssignment, ctx: SenderContext) -> J:
+            ctx.send_value(chained_assignment, attrgetter('_id'))
+            ctx.send_node(chained_assignment, attrgetter('_prefix'), PythonSender.send_space)
+            ctx.send_node(chained_assignment, attrgetter('_markers'), ctx.send_markers)
+            ctx.send_nodes(chained_assignment, attrgetter('_variables'), PythonSender.send_right_padded, lambda t: t.element.id)
+            ctx.send_node(chained_assignment, attrgetter('_assignment'), ctx.send_tree)
+            ctx.send_typed_value(chained_assignment, attrgetter('_type'))
+            return chained_assignment
+
         def visit_exception_type(self, exception_type: ExceptionType, ctx: SenderContext) -> J:
             ctx.send_value(exception_type, attrgetter('_id'))
             ctx.send_node(exception_type, attrgetter('_prefix'), PythonSender.send_space)

diff --git a/rewrite/rewrite/python/support_types.py b/rewrite/rewrite/python/support_types.py
@@ -32,6 +32,8 @@ class Location(Enum):
         BINARY_NEGATION = auto()
         BINARY_OPERATOR = auto()
         BINARY_PREFIX = auto()
+        CHAINED_ASSIGNMENT_PREFIX = auto()
+        CHAINED_ASSIGNMENT_VARIABLE_SUFFIX = auto()
         COLLECTION_LITERAL_ELEMENT_SUFFIX = auto()
         COLLECTION_LITERAL_PREFIX = auto()
         COMPILATION_UNIT_STATEMENT_PREFIX = auto()
@@ -81,7 +83,7 @@ class Location(Enum):
         TRAILING_ELSE_WRAPPER_PREFIX = auto()
         TYPE_HINTED_EXPRESSION_PREFIX = auto()
         TYPE_HINT_PREFIX = auto()
-        UNION_PREFIX = auto()
+        UNION_TYPE_PREFIX = auto()
         UNION_TYPES_SUFFIX = auto()
         VARIABLE_SCOPE_NAME_SUFFIX = auto()
         VARIABLE_SCOPE_PREFIX = auto()
@@ -96,6 +98,7 @@ class Location(Enum):
 
 class PyRightPadded:
     class Location(Enum):
+        CHAINED_ASSIGNMENT_VARIABLES = PySpace.Location.CHAINED_ASSIGNMENT_VARIABLE_SUFFIX
         ASSERT_STATEMENT_EXPRESSIONS = PySpace.Location.ASSERT_STATEMENT_EXPRESSION_SUFFIX
         COLLECTION_LITERAL_ELEMENT = PySpace.Location.COLLECTION_LITERAL_ELEMENT_SUFFIX
         COMPILATION_UNIT_STATEMENTS = PySpace.Location.COMPILATION_UNIT_STATEMENT_PREFIX
@@ -112,7 +115,7 @@ class Location(Enum):
         SLICE_STEP = PySpace.Location.SLICE_STEP_SUFFIX
         SLICE_STOP = PySpace.Location.SLICE_STOP_SUFFIX
         TOP_LEVEL_STATEMENT_SUFFIX = PySpace.Location.TOP_LEVEL_STATEMENT
-        UNION_TYPES = PySpace.Location.UNION_TYPES_SUFFIX
+        UNION_TYPE_TYPES = PySpace.Location.UNION_TYPES_SUFFIX
         VARIABLE_SCOPE_NAMES = PySpace.Location.VARIABLE_SCOPE_NAME_SUFFIX
         YIELD_EXPRESSIONS = PySpace.Location.YIELD_EXPRESSION_SUFFIX