From 25684e71310642ffd20b45eea9b5226a1fa809a5 Mon Sep 17 00:00:00 2001
From: Nikita Sobolev <mail@sobolevn.me>
Date: Wed, 13 Mar 2024 21:12:40 +0300
Subject: [PATCH] gh-100746: Improve `test_named_expressions.py` (#116713)

---
 Lib/test/test_named_expressions.py | 70 +++++++++++++++++++++++++++++-
 1 file changed, 68 insertions(+), 2 deletions(-)

diff --git a/Lib/test/test_named_expressions.py b/Lib/test/test_named_expressions.py
index f2017bdffcf..cf44080670d 100644
--- a/Lib/test/test_named_expressions.py
+++ b/Lib/test/test_named_expressions.py
@@ -298,6 +298,72 @@ class NamedExpressionInvalidTest(unittest.TestCase):
                 with self.assertRaisesRegex(SyntaxError, msg):
                     exec(f"lambda: {code}", {}) # Function scope
 
+    def test_named_expression_invalid_rebinding_dict_comprehension_iteration_variable(self):
+        cases = [
+            ("Key reuse", 'i', "{(i := 0): 1 for i in range(5)}"),
+            ("Value reuse", 'i', "{1: (i := 0) for i in range(5)}"),
+            ("Both reuse", 'i', "{(i := 0): (i := 0) for i in range(5)}"),
+            ("Nested reuse", 'j', "{{(j := 0): 1 for i in range(5)} for j in range(5)}"),
+            ("Reuse inner loop target", 'j', "{(j := 0): 1 for i in range(5) for j in range(5)}"),
+            ("Unpacking key reuse", 'i', "{(i := 0): 1 for i, j in {(0, 1)}}"),
+            ("Unpacking value reuse", 'i', "{1: (i := 0) for i, j in {(0, 1)}}"),
+            ("Reuse in loop condition", 'i', "{i+1: 1 for i in range(5) if (i := 0)}"),
+            ("Unreachable reuse", 'i', "{(False or (i:=0)): 1 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): 1 for a, (*b, c[d+e::f(g)], h.i) in j}"),
+            ("Complex expression: b", "b",
+                "{(b := 1): 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, {}) # 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_dict_comprehension_inner_loop(self):
+        cases = [
+            ("Inner reuse", 'j', "{i: 1 for i in range(5) if (j := 0) for j in range(5)}"),
+            ("Inner unpacking reuse", 'j', "{i: 1 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_dict_comprehension_iterable_expression(self):
+        cases = [
+            ("Top level", "{i: 1 for i in (i := range(5))}"),
+            ("Inside tuple", "{i: 1 for i in (2, 3, i := range(5))}"),
+            ("Inside list", "{i: 1 for i in [2, 3, i := range(5)]}"),
+            ("Different name", "{i: 1 for i in (j := range(5))}"),
+            ("Lambda expression", "{i: 1 for i in (lambda:(j := range(5)))()}"),
+            ("Inner loop", "{i: 1 for i in range(5) for j in (i := range(5))}"),
+            ("Nested comprehension", "{i: 1 for i in {j: 2 for j in (k := range(5))}}"),
+            ("Nested comprehension condition", "{i: 1 for i in {j: 2 for j in range(5) if (j := True)}}"),
+            ("Nested comprehension body", "{i: 1 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
+
     def test_named_expression_invalid_mangled_class_variables(self):
         code = """class Foo:
             def bar(self):
@@ -361,7 +427,7 @@ class NamedExpressionAssignmentTest(unittest.TestCase):
 
     def test_named_expression_assignment_10(self):
         if (match := 10) == 10:
-            pass
+            self.assertEqual(match, 10)
         else: self.fail("variable was not assigned using named expression")
 
     def test_named_expression_assignment_11(self):
@@ -403,7 +469,7 @@ class NamedExpressionAssignmentTest(unittest.TestCase):
 
     def test_named_expression_assignment_15(self):
         while a := False:
-            pass  # This will not run
+            self.fail("While body executed")  # This will not run
 
         self.assertEqual(a, False)