[3.9] bpo-42381: Allow walrus in set literals and set comprehensions (GH-23332) (GH-23333)

Currently walruses are not allowerd in set literals and set comprehensions:

>>> {y := 4, 4**2, 3**3}
  File "<stdin>", line 1
    {y := 4, 4**2, 3**3}
       ^
SyntaxError: invalid syntax

but they should be allowed as well per PEP 572.
(cherry picked from commit b0aba1f)

Co-authored-by: Pablo Galindo <[email protected]>

Author: pablogsal

Grammar/python.gram

@@ -301,7 +301,6 @@ block[asdl_seq*] (memo):
     | simple_stmt
     | invalid_block
 
-expressions_list[asdl_seq*]: a=','.star_expression+ [','] { a }
 star_expressions[expr_ty]:
     | a=star_expression b=(',' c=star_expression { c })+ [','] {
         _Py_Tuple(CHECK(_PyPegen_seq_insert_in_front(p, a, b)), Load, EXTRA) }
@@ -505,9 +504,9 @@ group[expr_ty]:
 genexp[expr_ty]:
     | '(' a=named_expression ~ b=for_if_clauses ')' { _Py_GeneratorExp(a, b, EXTRA) }
     | invalid_comprehension
-set[expr_ty]: '{' a=expressions_list '}' { _Py_Set(a, EXTRA) }
+set[expr_ty]: '{' a=star_named_expressions '}' { _Py_Set(a, EXTRA) }
 setcomp[expr_ty]:
-    | '{' a=expression ~ b=for_if_clauses '}' { _Py_SetComp(a, b, EXTRA) }
+    | '{' a=named_expression ~ b=for_if_clauses '}' { _Py_SetComp(a, b, EXTRA) }
     | invalid_comprehension
 dict[expr_ty]:
     | '{' a=[double_starred_kvpairs] '}' {

Lib/test/test_named_expressions.py

@@ -113,7 +113,7 @@ def test_named_expression_invalid_in_class_body(self):
             "assignment expression within a comprehension cannot be used in a class body"):
             exec(code, {}, {})
 
-    def test_named_expression_invalid_rebinding_comprehension_iteration_variable(self):
+    def test_named_expression_invalid_rebinding_list_comprehension_iteration_variable(self):
         cases = [
             ("Local reuse", 'i', "[i := 0 for i in range(5)]"),
             ("Nested reuse", 'j', "[[(j := 0) for i in range(5)] for j in range(5)]"),
@@ -130,7 +130,7 @@ def test_named_expression_invalid_rebinding_comprehension_iteration_variable(sel
                 with self.assertRaisesRegex(SyntaxError, msg):
                     exec(code, {}, {})
 
-    def test_named_expression_invalid_rebinding_comprehension_inner_loop(self):
+    def test_named_expression_invalid_rebinding_list_comprehension_inner_loop(self):
         cases = [
             ("Inner reuse", 'j', "[i for i in range(5) if (j := 0) for j in range(5)]"),
             ("Inner unpacking reuse", 'j', "[i for i in range(5) if (j := 0) for j, k in [(0, 1)]]"),
@@ -145,7 +145,7 @@ def test_named_expression_invalid_rebinding_comprehension_inner_loop(self):
                 with self.assertRaisesRegex(SyntaxError, msg):
                     exec(f"lambda: {code}", {}) # Function scope
 
-    def test_named_expression_invalid_comprehension_iterable_expression(self):
+    def test_named_expression_invalid_list_comprehension_iterable_expression(self):
         cases = [
             ("Top level", "[i for i in (i := range(5))]"),
             ("Inside tuple", "[i for i in (2, 3, i := range(5))]"),
@@ -167,6 +167,60 @@ def test_named_expression_invalid_comprehension_iterable_expression(self):
                 with self.assertRaisesRegex(SyntaxError, msg):
                     exec(f"lambda: {code}", {}) # Function scope
 
+    def test_named_expression_invalid_rebinding_set_comprehension_iteration_variable(self):
+        cases = [
+            ("Local reuse", 'i', "{i := 0 for i in range(5)}"),
+            ("Nested reuse", 'j', "{{(j := 0) for i in range(5)} for j in range(5)}"),
+            ("Reuse inner loop target", 'j', "{(j := 0) for i in range(5) for j in range(5)}"),
+            ("Unpacking reuse", 'i', "{i := 0 for i, j in {(0, 1)}}"),
+            ("Reuse in loop condition", 'i', "{i+1 for i in range(5) if (i := 0)}"),
+            ("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)}"),
+        ]
+        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, {}, {})
+
+    def test_named_expression_invalid_rebinding_set_comprehension_inner_loop(self):
+        cases = [
+            ("Inner reuse", 'j', "{i for i in range(5) if (j := 0) for j in range(5)}"),
+            ("Inner unpacking reuse", 'j', "{i for i in range(5) if (j := 0) for j, k in {(0, 1)}}"),
+        ]
+        for case, target, code in cases:
+            msg = f"comprehension inner loop cannot rebind assignment expression target '{target}'"
+            with self.subTest(case=case):
+                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_set_comprehension_iterable_expression(self):
+        cases = [
+            ("Top level", "{i for i in (i := range(5))}"),
+            ("Inside tuple", "{i for i in (2, 3, i := range(5))}"),
+            ("Inside list", "{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)}}"),
+        ]
+        msg = "assignment expression cannot be used in a comprehension iterable expression"
+        for case, code in cases:
+            with self.subTest(case=case):
+                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
+
 
 class NamedExpressionAssignmentTest(unittest.TestCase):
 

Misc/NEWS.d/next/Core and Builtins/2020-11-16-23-45-56.bpo-42381.G4AWxL.rst

@@ -0,0 +1,2 @@
+Allow assignment expressions in set literals and set comprehensions as per
+PEP 572. Patch by Pablo Galindo.