summaryrefslogtreecommitdiffstats
path: root/Lib/test
diff options
context:
space:
mode:
authorSerhiy Storchaka <storchaka@gmail.com>2020-02-12 10:18:59 (GMT)
committerGitHub <noreply@github.com>2020-02-12 10:18:59 (GMT)
commit8c579b1cc86053473eb052b76327279476740c9b (patch)
treebc99d5f52e1330500216512d1064c2f341b64d89 /Lib/test
parent0cc6b5e559b8303b18fdd56c2befd900fe7b5e35 (diff)
downloadcpython-8c579b1cc86053473eb052b76327279476740c9b.zip
cpython-8c579b1cc86053473eb052b76327279476740c9b.tar.gz
cpython-8c579b1cc86053473eb052b76327279476740c9b.tar.bz2
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`.
Diffstat (limited to 'Lib/test')
-rw-r--r--Lib/test/test_dictcomps.py17
-rw-r--r--Lib/test/test_genexps.py16
-rw-r--r--Lib/test/test_listcomps.py16
-rw-r--r--Lib/test/test_peepholer.py14
-rw-r--r--Lib/test/test_setcomps.py16
5 files changed, 79 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()
diff --git a/Lib/test/test_genexps.py b/Lib/test/test_genexps.py
index fd712bb..86e4e19 100644
--- a/Lib/test/test_genexps.py
+++ b/Lib/test/test_genexps.py
@@ -15,6 +15,22 @@ Test nesting with the inner expression dependent on the outer
>>> list((i,j) for i in range(4) for j in range(i) )
[(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
+Test the idiom for temporary variable assignment in comprehensions.
+
+ >>> list((j*j for i in range(4) for j in [i+1]))
+ [1, 4, 9, 16]
+ >>> list((j*k for i in range(4) for j in [i+1] for k in [j+1]))
+ [2, 6, 12, 20]
+ >>> list((j*k for i in range(4) for j, k in [(i+1, i+2)]))
+ [2, 6, 12, 20]
+
+Not assignment
+
+ >>> list((i*i for i in [*range(4)]))
+ [0, 1, 4, 9]
+ >>> list((i*i for i in (*range(4),)))
+ [0, 1, 4, 9]
+
Make sure the induction variable is not exposed
>>> i = 20
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index ddb169f..62b3319 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -16,6 +16,22 @@ Test nesting with the inner expression dependent on the outer
>>> [(i,j) for i in range(4) for j in range(i)]
[(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
+Test the idiom for temporary variable assignment in comprehensions.
+
+ >>> [j*j for i in range(4) for j in [i+1]]
+ [1, 4, 9, 16]
+ >>> [j*k for i in range(4) for j in [i+1] for k in [j+1]]
+ [2, 6, 12, 20]
+ >>> [j*k for i in range(4) for j, k in [(i+1, i+2)]]
+ [2, 6, 12, 20]
+
+Not assignment
+
+ >>> [i*i for i in [*range(4)]]
+ [0, 1, 4, 9]
+ >>> [i*i for i in (*range(4),)]
+ [0, 1, 4, 9]
+
Make sure the induction variable is not exposed
>>> i = 20
diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py
index 567e6a1..7913e91 100644
--- a/Lib/test/test_peepholer.py
+++ b/Lib/test/test_peepholer.py
@@ -495,6 +495,20 @@ class TestTranforms(BytecodeTestCase):
return 6
self.check_lnotab(f)
+ def test_assignment_idiom_in_comprehensions(self):
+ def listcomp():
+ return [y for x in a for y in [f(x)]]
+ self.assertEqual(count_instr_recursively(listcomp, 'FOR_ITER'), 1)
+ def setcomp():
+ return {y for x in a for y in [f(x)]}
+ self.assertEqual(count_instr_recursively(setcomp, 'FOR_ITER'), 1)
+ def dictcomp():
+ return {y: y for x in a for y in [f(x)]}
+ self.assertEqual(count_instr_recursively(dictcomp, 'FOR_ITER'), 1)
+ def genexpr():
+ return (y for x in a for y in [f(x)])
+ self.assertEqual(count_instr_recursively(genexpr, 'FOR_ITER'), 1)
+
class TestBuglets(unittest.TestCase):
diff --git a/Lib/test/test_setcomps.py b/Lib/test/test_setcomps.py
index fb7cde0..ecc4fff 100644
--- a/Lib/test/test_setcomps.py
+++ b/Lib/test/test_setcomps.py
@@ -21,6 +21,22 @@ Test nesting with the inner expression dependent on the outer
>>> list(sorted({(i,j) for i in range(4) for j in range(i)}))
[(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
+Test the idiom for temporary variable assignment in comprehensions.
+
+ >>> sorted({j*j for i in range(4) for j in [i+1]})
+ [1, 4, 9, 16]
+ >>> sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]})
+ [2, 6, 12, 20]
+ >>> sorted({j*k for i in range(4) for j, k in [(i+1, i+2)]})
+ [2, 6, 12, 20]
+
+Not assignment
+
+ >>> sorted({i*i for i in [*range(4)]})
+ [0, 1, 4, 9]
+ >>> sorted({i*i for i in (*range(4),)})
+ [0, 1, 4, 9]
+
Make sure the induction variable is not exposed
>>> i = 20