From 476989b083b906edf9bc9aaa70d1087a16f44d85 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Wed, 29 Jul 2020 21:56:10 +0300
Subject: [PATCH] Highlight assignment expressions used in iterables in
 comprehensions (PY-43619)

GitOrigin-RevId: a980602a3e90bbcfd5db03e828a5ad2a8a339569
---
 .../messages/PyBundle.properties              |  1 +
 .../validation/AssignTargetAnnotator.java     | 23 +++++++++++++++----
 .../assignmentExpressionInAnIterable.py       |  3 +++
 .../python/PythonHighlightingTest.java        |  5 ++++
 4 files changed, 28 insertions(+), 4 deletions(-)
 create mode 100644 python/testData/highlighting/assignmentExpressionInAnIterable.py

diff --git a/python/pluginResources/messages/PyBundle.properties b/python/pluginResources/messages/PyBundle.properties
index 8999fc20ed49..fa6a9086a071 100644
--- a/python/pluginResources/messages/PyBundle.properties
+++ b/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
diff --git a/python/src/com/jetbrains/python/validation/AssignTargetAnnotator.java b/python/src/com/jetbrains/python/validation/AssignTargetAnnotator.java
index 718a64ec5e8a..44031640f96a 100644
--- a/python/src/com/jetbrains/python/validation/AssignTargetAnnotator.java
+++ b/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();
     }
diff --git a/python/testData/highlighting/assignmentExpressionInAnIterable.py b/python/testData/highlighting/assignmentExpressionInAnIterable.py
new file mode 100644
index 000000000000..4bd9413ad49f
--- /dev/null
+++ b/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)]
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java b/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java
index 5d8e5e382962..09a9a301b041 100644
--- a/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java
+++ b/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();