bpo-32856: Optimize the assignment idiom in comprehensions. (GH-16814)

Now `for y in [expr]` in comprehensions is as fast as a simple
assignment `y = expr`.

1 files changed, 17 insertions, 0 deletions

diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py
index 927e310..16aa651 100644
--- a/Lib/test/test_dictcomps.py
+++ b/Lib/test/test_dictcomps.py
@@ -111,5 +111,22 @@ class DictComprehensionTest(unittest.TestCase):
         self.assertEqual(actual, expected)
         self.assertEqual(actual_calls, expected_calls)
 
+    def test_assignment_idiom_in_comprehensions(self):
+        expected = {1: 1, 2: 4, 3: 9, 4: 16}
+        actual = {j: j*j for i in range(4) for j in [i+1]}
+        self.assertEqual(actual, expected)
+        expected = {3: 2, 5: 6, 7: 12, 9: 20}
+        actual = {j+k: j*k for i in range(4) for j in [i+1] for k in [j+1]}
+        self.assertEqual(actual, expected)
+        expected = {3: 2, 5: 6, 7: 12, 9: 20}
+        actual = {j+k: j*k for i in range(4)  for j, k in [(i+1, i+2)]}
+        self.assertEqual(actual, expected)
+
+    def test_star_expression(self):
+        expected = {0: 0, 1: 1, 2: 4, 3: 9}
+        self.assertEqual({i: i*i for i in [*range(4)]}, expected)
+        self.assertEqual({i: i*i for i in (*range(4),)}, expected)
+
+
 if __name__ == "__main__":
     unittest.main()