diff options
author | Serhiy Storchaka <storchaka@gmail.com> | 2020-02-12 10:18:59 (GMT) |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-02-12 10:18:59 (GMT) |
commit | 8c579b1cc86053473eb052b76327279476740c9b (patch) | |
tree | bc99d5f52e1330500216512d1064c2f341b64d89 /Lib/test | |
parent | 0cc6b5e559b8303b18fdd56c2befd900fe7b5e35 (diff) | |
download | cpython-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.py | 17 | ||||
-rw-r--r-- | Lib/test/test_genexps.py | 16 | ||||
-rw-r--r-- | Lib/test/test_listcomps.py | 16 | ||||
-rw-r--r-- | Lib/test/test_peepholer.py | 14 | ||||
-rw-r--r-- | Lib/test/test_setcomps.py | 16 |
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 |