Highlight assignment expressions used in iterables in comprehensions (PY-43619)

GitOrigin-RevId: a980602a3e90bbcfd5db03e828a5ad2a8a339569

python/pluginResources/messages/PyBundle.properties

@@ -218,6 +218,7 @@ ANN.unparenthesized.assignment.expression.statement=Unparenthesized assignment e
 ANN.unparenthesized.assignment.expression.value=Unparenthesized assignment expressions are prohibited at the top level of the right hand side of an assignment statement
 ANN.assignment.expressions.within.a.comprehension.cannot.be.used.in.a.class.body=Assignment expressions within a comprehension cannot be used in a class body
 ANN.assignment.expression.as.a.target=Assignment expression cannot be used as a target here
+ANN.assignment.expression.in.an.iterable=Assignment expression cannot be used in a comprehension iterable
 ANN.ignore.errors.like.this=Ignore errors like this
 ANN.function.cannot.be.async=function \"{0}\" cannot be async
 ANN.python.does.not.support.yield.from.inside.async.functions=Python does not support 'yield from' inside async functions

python/src/com/jetbrains/python/validation/AssignTargetAnnotator.java

@@ -15,6 +15,7 @@
  */
 package com.jetbrains.python.validation;
 
+import com.intellij.codeInspection.util.InspectionMessage;
 import com.intellij.lang.annotation.HighlightSeverity;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
@@ -81,7 +82,7 @@ public class AssignTargetAnnotator extends PyAnnotator {
     PyExpression target = node.getForPart().getTarget();
     if (target != null) {
       target.accept(new ExprVisitor(Operation.For));
-      checkTargetIsNotAssignmentExpression(target);
+      checkNotAssignmentExpression(target, PyBundle.message("ANN.assignment.expression.as.a.target"));
     }
   }
 
@@ -115,13 +116,27 @@ public class AssignTargetAnnotator extends PyAnnotator {
 
   @Override
   public void visitPyComprehensionElement(@NotNull PyComprehensionElement node) {
-    node.getForComponents().forEach(it -> checkTargetIsNotAssignmentExpression(it.getIteratorVariable()));
+    final String targetMessage = PyBundle.message("ANN.assignment.expression.as.a.target");
+    final String iterableMessage = PyBundle.message("ANN.assignment.expression.in.an.iterable");
+
+    node.getForComponents().forEach(
+      it -> {
+        checkNotAssignmentExpression(it.getIteratorVariable(), targetMessage);
+        checkNoAssignmentExpressionAsChild(it.getIteratedList(), iterableMessage);
+      }
+    );
   }
 
-  private void checkTargetIsNotAssignmentExpression(@Nullable PyExpression expression) {
+  private void checkNoAssignmentExpressionAsChild(@Nullable PyExpression expression, @NotNull @InspectionMessage String message) {
+    PsiTreeUtil
+      .findChildrenOfType(expression, PyAssignmentExpression.class)
+      .forEach(it -> checkNotAssignmentExpression(it, message));
+  }
+
+  private void checkNotAssignmentExpression(@Nullable PyExpression expression, @NotNull @InspectionMessage String message) {
     if (PyPsiUtils.flattenParens(expression) instanceof PyAssignmentExpression) {
       getHolder()
-        .newAnnotation(HighlightSeverity.ERROR, PyBundle.message("ANN.assignment.expression.as.a.target"))
+        .newAnnotation(HighlightSeverity.ERROR, message)
         .range(expression)
         .create();
     }

python/testData/highlighting/assignmentExpressionInAnIterable.py

@@ -0,0 +1,3 @@
+[i for i in (<error descr="Assignment expression cannot be used in a comprehension iterable">x := (1, 2)</error>)]
+
+[[x for x in range(<error descr="Assignment expression cannot be used in a comprehension iterable">z := 10</error>)] for j in range(10)]

python/testSrc/com/jetbrains/python/PythonHighlightingTest.java

@@ -493,6 +493,11 @@ public class PythonHighlightingTest extends PyTestCase {
     doTest(LanguageLevel.getLatest(), false, false);
   }
 
+  // PY-43619
+  public void testAssignmentExpressionInAnIterable() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
   @NotNull
   private static EditorColorsScheme createTemporaryColorScheme() {
     EditorColorsManager manager = EditorColorsManager.getInstance();