[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 b0aba1fcdc3da952698d99aec2334faa79a8b68c)

Co-authored-by: Pablo Galindo <Pablogsal@gmail.com>

1 files changed, 57 insertions, 3 deletions

diff --git a/Lib/test/test_named_expressions.py b/Lib/test/test_named_expressions.py
index d614f32..2adcd4b 100644
--- a/Lib/test/test_named_expressions.py
+++ b/Lib/test/test_named_expressions.py
@@ -113,7 +113,7 @@ class NamedExpressionInvalidTest(unittest.TestCase):
             "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 @@ class NamedExpressionInvalidTest(unittest.TestCase):
                 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 @@ class NamedExpressionInvalidTest(unittest.TestCase):
                 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 @@ class NamedExpressionInvalidTest(unittest.TestCase):
                 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):