[Python-checkins] bpo-32856: Optimize the assignment idiom in comprehensions. (GH-16814)
Serhiy Storchaka
webhook-mailer at python.org
Wed Feb 12 05:19:08 EST 2020
https://github.com/python/cpython/commit/8c579b1cc86053473eb052b76327279476740c9b
commit: 8c579b1cc86053473eb052b76327279476740c9b
branch: master
author: Serhiy Storchaka <storchaka at gmail.com>
committer: GitHub <noreply at github.com>
date: 2020-02-12T12:18:59+02:00
summary:
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`.
files:
A Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst
M Doc/whatsnew/3.9.rst
M Lib/test/test_dictcomps.py
M Lib/test/test_genexps.py
M Lib/test/test_listcomps.py
M Lib/test/test_peepholer.py
M Lib/test/test_setcomps.py
M Python/compile.c
diff --git a/Doc/whatsnew/3.9.rst b/Doc/whatsnew/3.9.rst
index c0404101d7ce2..ec179845aee7c 100644
--- a/Doc/whatsnew/3.9.rst
+++ b/Doc/whatsnew/3.9.rst
@@ -315,6 +315,17 @@ case), and one used ``__VENV_NAME__`` instead.
Optimizations
=============
+* Optimized the idiom for assignment a temporary variable in comprehensions.
+ Now ``for y in [expr]`` in comprehensions is as fast as a simple assignment
+ ``y = expr``. For example:
+
+ sums = [s for s in [0] for x in data for s in [s + x]]
+
+ Unlike to the ``:=`` operator this idiom does not leak a variable to the
+ outer scope.
+
+ (Contributed by Serhiy Storchaka in :issue:`32856`.)
+
Build and C API Changes
=======================
diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py
index 927e3103e664b..16aa651b93c46 100644
--- a/Lib/test/test_dictcomps.py
+++ b/Lib/test/test_dictcomps.py
@@ -111,5 +111,22 @@ def add_call(pos, value):
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 fd712bb172d5d..86e4e195f55ec 100644
--- a/Lib/test/test_genexps.py
+++ b/Lib/test/test_genexps.py
@@ -15,6 +15,22 @@
>>> 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 ddb169fe58957..62b3319ad936d 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -16,6 +16,22 @@
>>> [(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 567e6a14361d1..7913e91e453ca 100644
--- a/Lib/test/test_peepholer.py
+++ b/Lib/test/test_peepholer.py
@@ -495,6 +495,20 @@ def f(x):
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 fb7cde03d7823..ecc4fffec0d84 100644
--- a/Lib/test/test_setcomps.py
+++ b/Lib/test/test_setcomps.py
@@ -21,6 +21,22 @@
>>> 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
diff --git a/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst b/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst
new file mode 100644
index 0000000000000..c1cd68f672712
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst
@@ -0,0 +1,3 @@
+Optimized the idiom for assignment a temporary variable in comprehensions.
+Now ``for y in [expr]`` in comprehensions is as fast as a simple assignment
+``y = expr``.
diff --git a/Python/compile.c b/Python/compile.c
index 04b8fe46e194d..bf8c8109d0758 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -212,11 +212,13 @@ static int compiler_set_qualname(struct compiler *);
static int compiler_sync_comprehension_generator(
struct compiler *c,
asdl_seq *generators, int gen_index,
+ int depth,
expr_ty elt, expr_ty val, int type);
static int compiler_async_comprehension_generator(
struct compiler *c,
asdl_seq *generators, int gen_index,
+ int depth,
expr_ty elt, expr_ty val, int type);
static PyCodeObject *assemble(struct compiler *, int addNone);
@@ -4343,22 +4345,24 @@ compiler_call_helper(struct compiler *c,
static int
compiler_comprehension_generator(struct compiler *c,
asdl_seq *generators, int gen_index,
+ int depth,
expr_ty elt, expr_ty val, int type)
{
comprehension_ty gen;
gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
if (gen->is_async) {
return compiler_async_comprehension_generator(
- c, generators, gen_index, elt, val, type);
+ c, generators, gen_index, depth, elt, val, type);
} else {
return compiler_sync_comprehension_generator(
- c, generators, gen_index, elt, val, type);
+ c, generators, gen_index, depth, elt, val, type);
}
}
static int
compiler_sync_comprehension_generator(struct compiler *c,
asdl_seq *generators, int gen_index,
+ int depth,
expr_ty elt, expr_ty val, int type)
{
/* generate code for the iterator, then each of the ifs,
@@ -4386,12 +4390,38 @@ compiler_sync_comprehension_generator(struct compiler *c,
}
else {
/* Sub-iter - calculate on the fly */
- VISIT(c, expr, gen->iter);
- ADDOP(c, GET_ITER);
+ /* Fast path for the temporary variable assignment idiom:
+ for y in [f(x)]
+ */
+ asdl_seq *elts;
+ switch (gen->iter->kind) {
+ case List_kind:
+ elts = gen->iter->v.List.elts;
+ break;
+ case Tuple_kind:
+ elts = gen->iter->v.Tuple.elts;
+ break;
+ default:
+ elts = NULL;
+ }
+ if (asdl_seq_LEN(elts) == 1) {
+ expr_ty elt = asdl_seq_GET(elts, 0);
+ if (elt->kind != Starred_kind) {
+ VISIT(c, expr, elt);
+ start = NULL;
+ }
+ }
+ if (start) {
+ VISIT(c, expr, gen->iter);
+ ADDOP(c, GET_ITER);
+ }
+ }
+ if (start) {
+ depth++;
+ compiler_use_next_block(c, start);
+ ADDOP_JREL(c, FOR_ITER, anchor);
+ NEXT_BLOCK(c);
}
- compiler_use_next_block(c, start);
- ADDOP_JREL(c, FOR_ITER, anchor);
- NEXT_BLOCK(c);
VISIT(c, expr, gen->target);
/* XXX this needs to be cleaned up...a lot! */
@@ -4405,7 +4435,7 @@ compiler_sync_comprehension_generator(struct compiler *c,
if (++gen_index < asdl_seq_LEN(generators))
if (!compiler_comprehension_generator(c,
- generators, gen_index,
+ generators, gen_index, depth,
elt, val, type))
return 0;
@@ -4420,18 +4450,18 @@ compiler_sync_comprehension_generator(struct compiler *c,
break;
case COMP_LISTCOMP:
VISIT(c, expr, elt);
- ADDOP_I(c, LIST_APPEND, gen_index + 1);
+ ADDOP_I(c, LIST_APPEND, depth + 1);
break;
case COMP_SETCOMP:
VISIT(c, expr, elt);
- ADDOP_I(c, SET_ADD, gen_index + 1);
+ ADDOP_I(c, SET_ADD, depth + 1);
break;
case COMP_DICTCOMP:
/* With '{k: v}', k is evaluated before v, so we do
the same. */
VISIT(c, expr, elt);
VISIT(c, expr, val);
- ADDOP_I(c, MAP_ADD, gen_index + 1);
+ ADDOP_I(c, MAP_ADD, depth + 1);
break;
default:
return 0;
@@ -4440,8 +4470,10 @@ compiler_sync_comprehension_generator(struct compiler *c,
compiler_use_next_block(c, skip);
}
compiler_use_next_block(c, if_cleanup);
- ADDOP_JABS(c, JUMP_ABSOLUTE, start);
- compiler_use_next_block(c, anchor);
+ if (start) {
+ ADDOP_JABS(c, JUMP_ABSOLUTE, start);
+ compiler_use_next_block(c, anchor);
+ }
return 1;
}
@@ -4449,6 +4481,7 @@ compiler_sync_comprehension_generator(struct compiler *c,
static int
compiler_async_comprehension_generator(struct compiler *c,
asdl_seq *generators, int gen_index,
+ int depth,
expr_ty elt, expr_ty val, int type)
{
comprehension_ty gen;
@@ -4492,9 +4525,10 @@ compiler_async_comprehension_generator(struct compiler *c,
NEXT_BLOCK(c);
}
+ depth++;
if (++gen_index < asdl_seq_LEN(generators))
if (!compiler_comprehension_generator(c,
- generators, gen_index,
+ generators, gen_index, depth,
elt, val, type))
return 0;
@@ -4509,18 +4543,18 @@ compiler_async_comprehension_generator(struct compiler *c,
break;
case COMP_LISTCOMP:
VISIT(c, expr, elt);
- ADDOP_I(c, LIST_APPEND, gen_index + 1);
+ ADDOP_I(c, LIST_APPEND, depth + 1);
break;
case COMP_SETCOMP:
VISIT(c, expr, elt);
- ADDOP_I(c, SET_ADD, gen_index + 1);
+ ADDOP_I(c, SET_ADD, depth + 1);
break;
case COMP_DICTCOMP:
/* With '{k: v}', k is evaluated before v, so we do
the same. */
VISIT(c, expr, elt);
VISIT(c, expr, val);
- ADDOP_I(c, MAP_ADD, gen_index + 1);
+ ADDOP_I(c, MAP_ADD, depth + 1);
break;
default:
return 0;
@@ -4583,7 +4617,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
ADDOP_I(c, op, 0);
}
- if (!compiler_comprehension_generator(c, generators, 0, elt,
+ if (!compiler_comprehension_generator(c, generators, 0, 0, elt,
val, type))
goto error_in_scope;
More information about the Python-checkins
mailing list