gh-87447: Fix walrus comprehension rebind checking (#100581)

Co-authored-by: Pablo Galindo Salgado <Pablogsal@gmail.com>
Co-authored-by: Shantanu <12621235+hauntsaninja@users.noreply.github.com>

4 files changed, 92 insertions, 3 deletions

diff --git a/Doc/whatsnew/3.12.rst b/Doc/whatsnew/3.12.rst
index b882bb6..7d318ca 100644
--- a/Doc/whatsnew/3.12.rst
+++ b/Doc/whatsnew/3.12.rst
@@ -182,6 +182,13 @@ Other Language Changes
   arguments of any type instead of just :class:`bool` and :class:`int`.
   (Contributed by Serhiy Storchaka in :gh:`60203`.)
 
+* Variables used in the target part of comprehensions that are not stored to
+  can now be used in assignment expressions (``:=``).
+  For example, in ``[(b := 1) for a, b.prop in some_iter]``, the assignment to
+  ``b`` is now allowed. Note that assigning to variables stored to in the target
+  part of comprehensions (like ``a``) is still disallowed, as per :pep:`572`.
+  (Contributed by Nikita Sobolev in :gh:`100581`.)
+
 
 New Modules
 ===========
diff --git a/Lib/test/test_named_expressions.py b/Lib/test/test_named_expressions.py
index 20ac2e6..7b2fa84 100644
--- a/Lib/test/test_named_expressions.py
+++ b/Lib/test/test_named_expressions.py
@@ -114,6 +114,69 @@ class NamedExpressionInvalidTest(unittest.TestCase):
             "assignment expression within a comprehension cannot be used in a class body"):
             exec(code, {}, {})
 
+    def test_named_expression_valid_rebinding_iteration_variable(self):
+        # This test covers that we can reassign variables
+        # that are not directly assigned in the
+        # iterable part of a comprehension.
+        cases = [
+            # Regression tests from https://github.com/python/cpython/issues/87447
+            ("Complex expression: c",
+                "{0}(c := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+            ("Complex expression: d",
+                "{0}(d := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+            ("Complex expression: e",
+                "{0}(e := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+            ("Complex expression: f",
+                "{0}(f := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+            ("Complex expression: g",
+                "{0}(g := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+            ("Complex expression: h",
+                "{0}(h := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+            ("Complex expression: i",
+                "{0}(i := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+            ("Complex expression: j",
+                "{0}(j := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+        ]
+        for test_case, code in cases:
+            for lpar, rpar in [('(', ')'), ('[', ']'), ('{', '}')]:
+                code = code.format(lpar, rpar)
+                with self.subTest(case=test_case, lpar=lpar, rpar=rpar):
+                    # Names used in snippets are not defined,
+                    # but we are fine with it: just must not be a SyntaxError.
+                    # Names used in snippets are not defined,
+                    # but we are fine with it: just must not be a SyntaxError.
+                    with self.assertRaises(NameError):
+                        exec(code, {}) # Module scope
+                    with self.assertRaises(NameError):
+                        exec(code, {}, {}) # Class scope
+                    exec(f"lambda: {code}", {}) # Function scope
+
+    def test_named_expression_invalid_rebinding_iteration_variable(self):
+        # This test covers that we cannot reassign variables
+        # that are directly assigned in the iterable part of a comprehension.
+        cases = [
+            # Regression tests from https://github.com/python/cpython/issues/87447
+            ("Complex expression: a", "a",
+                "{0}(a := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+            ("Complex expression: b", "b",
+                "{0}(b := 1) for a, (*b, c[d+e::f(g)], h.i) in j{1}"),
+        ]
+        for test_case, target, code in cases:
+            msg = f"assignment expression cannot rebind comprehension iteration variable '{target}'"
+            for lpar, rpar in [('(', ')'), ('[', ']'), ('{', '}')]:
+                code = code.format(lpar, rpar)
+                with self.subTest(case=test_case, lpar=lpar, rpar=rpar):
+                    # Names used in snippets are not defined,
+                    # but we are fine with it: just must not be a SyntaxError.
+                    # Names used in snippets are not defined,
+                    # but we are fine with it: just must not be a SyntaxError.
+                    with self.assertRaisesRegex(SyntaxError, msg):
+                        exec(code, {}) # Module scope
+                    with self.assertRaisesRegex(SyntaxError, msg):
+                        exec(code, {}, {}) # Class scope
+                    with self.assertRaisesRegex(SyntaxError, msg):
+                        exec(f"lambda: {code}", {}) # Function scope
+
     def test_named_expression_invalid_rebinding_list_comprehension_iteration_variable(self):
         cases = [
             ("Local reuse", 'i', "[i := 0 for i in range(5)]"),
@@ -129,7 +192,11 @@ class NamedExpressionInvalidTest(unittest.TestCase):
             msg = f"assignment expression cannot rebind comprehension iteration variable '{target}'"
             with self.subTest(case=case):
                 with self.assertRaisesRegex(SyntaxError, msg):
-                    exec(code, {}, {})
+                    exec(code, {}) # Module scope
+                with self.assertRaisesRegex(SyntaxError, msg):
+                    exec(code, {}, {}) # Class scope
+                with self.assertRaisesRegex(SyntaxError, msg):
+                    exec(f"lambda: {code}", {}) # Function scope
 
     def test_named_expression_invalid_rebinding_list_comprehension_inner_loop(self):
         cases = [
@@ -178,12 +245,21 @@ class NamedExpressionInvalidTest(unittest.TestCase):
             ("Unreachable reuse", 'i', "{False or (i:=0) for i in range(5)}"),
             ("Unreachable nested reuse", 'i',
                 "{(i, j) for i in range(5) for j in range(5) if True or (i:=10)}"),
+            # Regression tests from https://github.com/python/cpython/issues/87447
+            ("Complex expression: a", "a",
+                "{(a := 1) for a, (*b, c[d+e::f(g)], h.i) in j}"),
+            ("Complex expression: b", "b",
+                "{(b := 1) for a, (*b, c[d+e::f(g)], h.i) in j}"),
         ]
         for case, target, code in cases:
             msg = f"assignment expression cannot rebind comprehension iteration variable '{target}'"
             with self.subTest(case=case):
                 with self.assertRaisesRegex(SyntaxError, msg):
-                    exec(code, {}, {})
+                    exec(code, {}) # Module scope
+                with self.assertRaisesRegex(SyntaxError, msg):
+                    exec(code, {}, {}) # Class scope
+                with self.assertRaisesRegex(SyntaxError, msg):
+                    exec(f"lambda: {code}", {}) # Function scope
 
     def test_named_expression_invalid_rebinding_set_comprehension_inner_loop(self):
         cases = [
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-12-28-15-02-53.gh-issue-87447.7-aekA.rst b/Misc/NEWS.d/next/Core and Builtins/2022-12-28-15-02-53.gh-issue-87447.7-aekA.rst
new file mode 100644
index 0000000..af60acf
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-12-28-15-02-53.gh-issue-87447.7-aekA.rst
@@ -0,0 +1,5 @@
+Fix :exc:`SyntaxError` on comprehension rebind checking with names that are
+not actually redefined.
+
+Now reassigning ``b`` in ``[(b := 1) for a, b.prop in some_iter]`` is allowed.
+Reassigning ``a`` is still disallowed as per :pep:`572`.
diff --git a/Python/symtable.c b/Python/symtable.c
index 3c13018..89a2bc4 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -1488,7 +1488,8 @@ symtable_extend_namedexpr_scope(struct symtable *st, expr_ty e)
          */
         if (ste->ste_comprehension) {
             long target_in_scope = _PyST_GetSymbol(ste, target_name);
-            if (target_in_scope & DEF_COMP_ITER) {
+            if ((target_in_scope & DEF_COMP_ITER) &&
+                (target_in_scope & DEF_LOCAL)) {
                 PyErr_Format(PyExc_SyntaxError, NAMED_EXPR_COMP_CONFLICT, target_name);
                 PyErr_RangedSyntaxLocationObject(st->st_filename,
                                                   e->lineno,