Update scopes processing for assignment expressions (PEP 572) (PY-33886)

From the PEP:
an assignment expression occurring in
a list, set or dict comprehension or in a generator expression
binds the target in the containing scope.

GitOrigin-RevId: 88d701a8da298a9b0bdd8c7ac4d4940a5a22f9db

python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java

@@ -183,6 +183,8 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
   }
 
   private static boolean isInnerComprehension(PsiElement referenceElement, PsiElement definition) {
+    if (definition != null && definition.getParent() instanceof PyAssignmentExpression) return false;
+
     final PyComprehensionElement definitionComprehension = PsiTreeUtil.getParentOfType(definition, PyComprehensionElement.class);
     if (definitionComprehension != null && PyUtil.isOwnScopeComprehension(definitionComprehension)) {
       final PyComprehensionElement elementComprehension = PsiTreeUtil.getParentOfType(referenceElement, PyComprehensionElement.class);

python/testData/resolve/AssignmentExpressionsAndOuterVar.py

@@ -0,0 +1,4 @@
+total = 0
+partial_sums = [total := total + v for v in values]
+print("Total:", total)
+                 <ref>

python/testSrc/com/jetbrains/python/PyResolveTest.java

@@ -1,6 +1,7 @@
 // Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.jetbrains.python;
 
+import com.intellij.openapi.util.Ref;
 import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.testFramework.UsefulTestCase;
@@ -16,6 +17,8 @@ import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.types.PyClassTypeImpl;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.pyi.PyiUtil;
+import org.intellij.lang.annotations.Language;
+import org.jetbrains.annotations.NotNull;
 
 public class PyResolveTest extends PyResolveTestCase {
   @Override
@@ -41,6 +44,22 @@ public class PyResolveTest extends PyResolveTestCase {
     return ((PsiPolyVariantReference)ref).multiResolve(false);
   }
 
+  private <T extends PsiNamedElement> T assertResolvesTo(@Language("TEXT") @NotNull String text,
+                                                         @NotNull Class<T> cls,
+                                                         @NotNull String name) {
+    final Ref<T> result = new Ref<>();
+
+    runWithLanguageLevel(
+      LanguageLevel.getLatest(),
+      () -> {
+        myFixture.configureByText(PythonFileType.INSTANCE, text);
+        result.set(assertResolveResult(PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve(), cls, name));
+      }
+    );
+
+    return result.get();
+  }
+
   public void testClass() {
     assertResolvesTo(PyClass.class, "Test");
   }
@@ -1350,4 +1369,86 @@ public class PyResolveTest extends PyResolveTestCase {
     assertNotNull(function);
     assertFalse(PyiUtil.isOverload(function, context));
   }
+
+  // PY-33886
+  public void testAssignmentExpressions() {
+    assertResolvesTo(
+      "if a := b:\n" +
+      "    print(a)\n" +
+      "          <ref>",
+      PyTargetExpression.class,
+      "a"
+    );
+
+    assertResolvesTo(
+      "[y := 2, y**2]\n" +
+      "         <ref>",
+      PyTargetExpression.class,
+      "y"
+    );
+
+    assertResolvesTo(
+      "[y for x in data if (y := f(x))]\n" +
+      " <ref>",
+      PyTargetExpression.class,
+      "y"
+    );
+
+    assertResolvesTo(
+      "len(lines := [])\n" +
+      "print(lines)\n" +
+      "       <ref>",
+      PyTargetExpression.class,
+      "lines"
+    );
+  }
+
+  // PY-33886
+  public void testAssignmentExpressionGoesToOuterScope() {
+    assertResolvesTo(
+      "if any({(comment := line).startswith('#') for line in lines}):\n" +
+      "    print(\"First comment:\", comment)\n" +
+      "                             <ref>",
+      PyTargetExpression.class,
+      "comment"
+    );
+
+    assertResolvesTo(
+      "if all((nonblank := line).strip() == '' for line in lines):\n" +
+      "    pass\n" +
+      "else:\n" +
+      "    print(\"First non-blank line:\", nonblank)\n" +
+      "                                     <ref>",
+      PyTargetExpression.class,
+      "nonblank"
+    );
+  }
+
+  // PY-33886
+  public void testAssignmentExpressionGoesFromOuterScope() {
+    assertResolvesTo(
+      "[[x * y for x in range(5)] for i in range(5) if (y := i)]\n" +
+      "      <ref>",
+      PyTargetExpression.class,
+      "y"
+    );
+  }
+
+  // PY-33886
+  public void testAssignmentExpressionsAndOuterVar() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON38,
+      () -> {
+        final ResolveResult[] results = multiResolve();
+
+        final PsiElement first = results[0].getElement();
+        assertResolveResult(first, PyTargetExpression.class, "total");
+        assertInstanceOf(first.getParent(), PyAssignmentStatement.class);
+
+        final PsiElement second = results[1].getElement();
+        assertResolveResult(second, PyTargetExpression.class, "total");
+        assertInstanceOf(second.getParent(), PyAssignmentExpression.class);
+      }
+    );
+  }
 }