From b600a9bebf02dad1a11dffb2912f027276e1887c Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Wed, 24 Apr 2019 23:38:42 +0300
Subject: [PATCH] Suppress "Unbound local variable" for targets defined in
 assignment expressions inside comprehensions under conditions (PEP 572)
 (PY-33886)

GitOrigin-RevId: ef7edfd58d6c36037903c54ed2c832696c4ff0b6
---
 .../PyUnboundLocalVariableInspection.java     | 13 ++++++-
 .../PyUnboundLocalVariableInspectionTest.java | 34 +++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/python/src/com/jetbrains/python/inspections/PyUnboundLocalVariableInspection.java b/python/src/com/jetbrains/python/inspections/PyUnboundLocalVariableInspection.java
index adaece1caf4c..aa14acc58655 100644
--- a/python/src/com/jetbrains/python/inspections/PyUnboundLocalVariableInspection.java
+++ b/python/src/com/jetbrains/python/inspections/PyUnboundLocalVariableInspection.java
@@ -129,7 +129,7 @@ public class PyUnboundLocalVariableInspection extends PyInspection {
             return;
           }
         }
-        if (resolvedUnderWithStatement(node, resolved)) {
+        if (resolvedUnderWithStatement(node, resolved) || resolvedUnderAssignmentExpressionAndCondition(node, resolved)) {
           return;
         }
         if (PyUnreachableCodeInspection.hasAnyInterruptedControlFlowPaths(node)) {
@@ -159,6 +159,17 @@ public class PyUnboundLocalVariableInspection extends PyInspection {
              PsiTreeUtil.getParentOfType(resolved, PyWithStatement.class, true, ScopeOwner.class) != null;
     }
 
+    private static boolean resolvedUnderAssignmentExpressionAndCondition(@NotNull PyReferenceExpression node,
+                                                                         @Nullable PsiElement resolved) {
+      return resolved instanceof PyTargetExpression &&
+             PyUtil.inSameFile(node, resolved) &&
+             resolved.getParent() instanceof PyAssignmentExpression &&
+             PsiTreeUtil.getParentOfType(
+               PsiTreeUtil.getParentOfType(resolved.getParent(), PyComprehensionElement.class, true, PyConditionalStatementPart.class),
+               PyConditionalStatementPart.class
+             ) != null;
+    }
+
     private static boolean isFirstUnboundRead(@NotNull PyReferenceExpression node, @NotNull ScopeOwner owner) {
       final String nodeName = node.getReferencedName();
       final Scope scope = ControlFlowCache.getScope(owner);
diff --git a/python/testSrc/com/jetbrains/python/inspections/PyUnboundLocalVariableInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/PyUnboundLocalVariableInspectionTest.java
index 30eae2ec1521..bdb054d352be 100644
--- a/python/testSrc/com/jetbrains/python/inspections/PyUnboundLocalVariableInspectionTest.java
+++ b/python/testSrc/com/jetbrains/python/inspections/PyUnboundLocalVariableInspectionTest.java
@@ -290,6 +290,40 @@ public class PyUnboundLocalVariableInspectionTest extends PyInspectionTestCase {
                  "    print(val)");
   }
 
+  // PY-33886
+  public void testAssignmentExpressions() {
+    runWithLanguageLevel(
+      LanguageLevel.getLatest(),
+      () -> doTestByText(
+        "def foo():\n" +
+        "    if any((comment := line).startswith('#') for line in lines):\n" +
+        "        print(\"First comment:\", comment)\n" +
+        "    else:\n" +
+        "        print(\"There are no comments\")\n" +
+        "\n" +
+        "    if all((nonblank := line).strip() == '' for line in lines):\n" +
+        "        print(\"All lines are blank\")\n" +
+        "    else:\n" +
+        "        print(\"First non-blank line:\", nonblank)\n" +
+        "\n" +
+        "\n" +
+        "def bar():\n" +
+        "    [(comment := line).startswith('#') for line in lines]\n" +
+        "    print(<warning descr=\"Local variable 'comment' might be referenced before assignment\">comment</warning>)\n" +
+        "\n" +
+        "\n" +
+        "def baz():\n" +
+        "    while (line := input()) and any((first_digit := c).isdigit() for c in line):\n" +
+        "        print(line, first_digit)\n" +
+        "\n" +
+        "\n" +
+        "def more():\n" +
+        "    if (x := True) or (y := 'spam'):\n" +
+        "        print(<warning descr=\"Local variable 'y' might be referenced before assignment\">y</warning>)\n"
+      )
+    );
+  }
+
   @NotNull
   @Override
   protected Class<? extends PyInspection> getInspectionClass() {