traverseda
/
cpython
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

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`.
This commit is contained in:
Serhiy Storchaka 2020-02-12 12:18:59 +02:00 committed by GitHub
parent 0cc6b5e559
commit 8c579b1cc8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 145 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
Doc/whatsnew/3.9.rst
View File

@ -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
=======================

17
Lib/test/test_dictcomps.py
View File

@ -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()

16
Lib/test/test_genexps.py
View File

@ -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

16
Lib/test/test_listcomps.py
View File

@ -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

14
Lib/test/test_peepholer.py
View File

@ -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):

16
Lib/test/test_setcomps.py
View File

@ -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

3
Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst Normal file
View File

@ -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``.

56
Python/compile.c
View File

@ -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 @@ ex_call:
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 */
/* 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);
}
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);
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;