Close nested comprehension and lambda expression loopholes

python · ncoghlan · Aug 25, 2019 · Aug 5, 2019 · Aug 5, 2019 · Aug 5, 2019
commit c0f2b5bfa7f3b0318493a5d663ac59cf3f705df0
diff --git a/Lib/test/test_named_expressions.py b/Lib/test/test_named_expressions.py
@@ -123,8 +123,6 @@ def test_named_expression_invalid_rebinding_comprehension_iteration_variable(sel
             ("Unreachable reuse", "[False or (i:=0) for i in range(5)]"),
             ("Unreachable nested reuse",
                 "[(i, j) for i in range(5) for j in range(5) if True or (i:=10)]"),
-            ("Nested comprehension condition", "[i for i in [j for j in range(5) if (j := True)]]"),
-            ("Nested comprehension body", "[i for i in [(j := True) for j in range(5)]]"),
         ]
         for case, code in cases:
             with self.subTest(case=case):
@@ -148,8 +146,11 @@ def test_named_expression_invalid_comprehension_iterator_expression(self):
             ("Top level", "[i for i in (i := range(5))]"),
             ("Inside container", "[i for i in (2, 3, i := range(5))]"),
             ("Different name", "[i for i in (j := range(5))]"),
+            ("Lambda expression", "[i for i in (lambda:(j := range(5)))()]"),
             ("Inner loop", "[i for i in range(5) for j in (i := range(5))]"),
             ("Nested comprehension", "[i for i in [j for j in (k := range(5))]]"),
+            ("Nested comprehension condition", "[i for i in [j for j in range(5) if (j := True)]]"),
+            ("Nested comprehension body", "[i for i in [(j := True) for j in range(5)]]"),
         ]
         for case, code in cases:
             with self.subTest(case=case):

diff --git a/Python/symtable.c b/Python/symtable.c
@@ -973,6 +973,13 @@ symtable_enter_block(struct symtable *st, identifier name, _Py_block_ty block,
         return 0;
     }
     prev = st->st_cur;
+    /* bpo-37757: For now, disallow *all* assignment expressions in the
+     * outermost iterator expression of a comprehension, even those inside
+     * a nested comprehension or a lambda expression.
+     */
+    if (prev) {
+        ste->ste_comp_iter_expr = prev->ste_comp_iter_expr;
+    }
     /* The entry is owned by the stack. Borrow it for st_cur. */
     Py_DECREF(ste);
     st->st_cur = ste;