assignment expression in comprehension should target outer scope (#698)

Co-authored-by: Yann Sartori <sartori2@llnl.gov>

2 files changed, 33 insertions, 1 deletions

diff --git a/pyflakes/checker.py b/pyflakes/checker.py
index 7f33d6a..0c3f66e 100644
--- a/pyflakes/checker.py
+++ b/pyflakes/checker.py
@@ -540,6 +540,12 @@ class Assignment(Binding):
     """
 
 
+class NamedExprAssignment(Assignment):
+    """
+    Represents binding a name with an assignment expression.
+    """
+
+
 class Annotation(Binding):
     """
     Represents binding a name to a type without an associated value.
@@ -1159,7 +1165,14 @@ class Checker(object):
         # don't treat annotations as assignments if there is an existing value
         # in scope
         if value.name not in self.scope or not isinstance(value, Annotation):
-            self.scope[value.name] = value
+            cur_scope_pos = -1
+            # As per PEP 572, use scope in which outermost generator is defined
+            while (
+                isinstance(value, NamedExprAssignment) and
+                isinstance(self.scopeStack[cur_scope_pos], GeneratorScope)
+            ):
+                cur_scope_pos -= 1
+            self.scopeStack[cur_scope_pos][value.name] = value
 
     def _unknown_handler(self, node):
         # this environment variable configures whether to error on unknown
@@ -1302,6 +1315,8 @@ class Checker(object):
             binding = ExportBinding(name, node._pyflakes_parent, self.scope)
         elif PY2 and isinstance(getattr(node, 'ctx', None), ast.Param):
             binding = Argument(name, self.getScopeNode(node))
+        elif PY38_PLUS and isinstance(parent_stmt, ast.NamedExpr):
+            binding = NamedExprAssignment(name, node)
         else:
             binding = Assignment(name, node)
         self.addBinding(node, binding)
diff --git a/pyflakes/test/test_other.py b/pyflakes/test/test_other.py
index 68813bd..efbc75d 100644
--- a/pyflakes/test/test_other.py
+++ b/pyflakes/test/test_other.py
@@ -1772,6 +1772,23 @@ class TestUnusedAssignment(TestCase):
         print(x)
         ''')
 
+    @skipIf(version_info < (3, 8), 'new in Python 3.8')
+    def test_assign_expr_generator_scope(self):
+        """Test assignment expressions in generator expressions."""
+        self.flakes('''
+        if (any((y := x[0]) for x in [[True]])):
+            print(y)
+        ''')
+
+    @skipIf(version_info < (3, 8), 'new in Python 3.8')
+    def test_assign_expr_nested(self):
+        """Test assignment expressions in nested expressions."""
+        self.flakes('''
+        if ([(y:=x) for x in range(4) if [(z:=q) for q in range(4)]]):
+            print(y)
+            print(z)
+        ''')
+
 
 class TestStringFormatting(TestCase):