From aa1da541ded7dfef575e5e7099c8238513f20c84 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Wed, 24 Apr 2019 23:10:49 +0300
Subject: [PATCH] Highlight assignment expressions within a comprehension in a
 class body as errors (PEP 572) (PY-33886)

GitOrigin-RevId: 1b63427a4c5e4927be07fdc4bcec91bd952e26c7
---
 .../python/validation/AssignTargetAnnotator.java      | 11 +++++++++++
 .../highlighting/invalidAssignmentExpressions.py      |  5 ++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/python/src/com/jetbrains/python/validation/AssignTargetAnnotator.java b/python/src/com/jetbrains/python/validation/AssignTargetAnnotator.java
index 7b80b5cdea16..b55fea30c1f2 100644
--- a/python/src/com/jetbrains/python/validation/AssignTargetAnnotator.java
+++ b/python/src/com/jetbrains/python/validation/AssignTargetAnnotator.java
@@ -16,7 +16,10 @@
 package com.jetbrains.python.validation;
 
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.PyNames;
+import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.sdk.PythonSdkType;
 import org.jetbrains.annotations.NotNull;
@@ -104,6 +107,14 @@ public class AssignTargetAnnotator extends PyAnnotator {
     errorOnUnparenthesizedAssignmentExpression(node.getValue(), "as annotations for arguments, return values and assignments");
   }
 
+  @Override
+  public void visitPyAssignmentExpression(PyAssignmentExpression node) {
+    final PyComprehensionElement comprehensionElement = PsiTreeUtil.getParentOfType(node, PyComprehensionElement.class, true, ScopeOwner.class);
+    if (ScopeUtil.getScopeOwner(comprehensionElement) instanceof PyClass) {
+      getHolder().createErrorAnnotation(node, "Assignment expressions within a comprehension cannot be used in a class body");
+    }
+  }
+
   private void errorOnUnparenthesizedAssignmentExpression(@Nullable PyExpression expression, @NotNull String suffix) {
     if (expression instanceof PyAssignmentExpression) {
       getHolder().createErrorAnnotation(expression, "Unparenthesized assignment expressions are prohibited " + suffix);
diff --git a/python/testData/highlighting/invalidAssignmentExpressions.py b/python/testData/highlighting/invalidAssignmentExpressions.py
index 403617f85532..1be4714b88d8 100644
--- a/python/testData/highlighting/invalidAssignmentExpressions.py
+++ b/python/testData/highlighting/invalidAssignmentExpressions.py
@@ -20,4 +20,7 @@ def foo(answer: (p := 42) = 5):  # Valid, but probably never useful
 (lambda: <error descr="Unparenthesized assignment expressions are prohibited at the top level of a lambda function">x := 1</error>) # INVALID
 lambda: (x := 1) # Valid, but unlikely to be useful
 (x := lambda: 1) # Valid
-lambda line: (m := re.match(pattern, line)) and m.group(1) # Valid
\ No newline at end of file
+lambda line: (m := re.match(pattern, line)) and m.group(1) # Valid
+
+class A:
+    [<error descr="Assignment expressions within a comprehension cannot be used in a class body">y := i</error> for i in range(2)]
\ No newline at end of file