diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py index 0873071883439e..23812006dde246 100644 --- a/Lib/test/test_dictcomps.py +++ b/Lib/test/test_dictcomps.py @@ -81,6 +81,17 @@ def test_illegal_assignment(self): compile("{x: y for y, x in ((1, 2), (3, 4))} += 5", "", "exec") + 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) + if __name__ == "__main__": unittest.main() diff --git a/Lib/test/test_genexps.py b/Lib/test/test_genexps.py index fb531d6d472b6f..5305204900f361 100644 --- a/Lib/test/test_genexps.py +++ b/Lib/test/test_genexps.py @@ -15,6 +15,15 @@ >>> 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] + 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 ddb169fe58957c..da441da7d28b07 100644 --- a/Lib/test/test_listcomps.py +++ b/Lib/test/test_listcomps.py @@ -16,6 +16,15 @@ >>> [(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] + 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 794d104d5919bd..bd077102103d81 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -335,6 +335,20 @@ def forloop(): pass self.assertEqual(count_instr_recursively(forloop, 'BUILD_LIST'), 0) + 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 fb7cde03d78236..ad954b04259b86 100644 --- a/Lib/test/test_setcomps.py +++ b/Lib/test/test_setcomps.py @@ -21,6 +21,15 @@ >>> 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. + + >>> list(sorted({j*j for i in range(4) for j in [i+1]})) + [1, 4, 9, 16] + >>> list(sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]})) + [2, 6, 12, 20] + >>> list(sorted({j*k for i in range(4) for j, k in [(i+1, i+2)]})) + [2, 6, 12, 20] + 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 00000000000000..4233a919807a28 --- /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 so fast as a simple assignment +``y = expr``. diff --git a/Python/compile.c b/Python/compile.c index 3528670ef67503..6271e2ad1efa2f 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -200,11 +200,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); @@ -3875,22 +3877,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, @@ -3918,12 +3922,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! */ @@ -3937,7 +3967,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; @@ -3952,18 +3982,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 'd[k] = v', v is evaluated before k, so we do the same. */ VISIT(c, expr, val); VISIT(c, expr, elt); - ADDOP_I(c, MAP_ADD, gen_index + 1); + ADDOP_I(c, MAP_ADD, depth + 1); break; default: return 0; @@ -3972,8 +4002,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; } @@ -3981,6 +4013,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; @@ -4024,9 +4057,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; @@ -4041,18 +4075,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 'd[k] = v', v is evaluated before k, so we do the same. */ VISIT(c, expr, val); VISIT(c, expr, elt); - ADDOP_I(c, MAP_ADD, gen_index + 1); + ADDOP_I(c, MAP_ADD, depth + 1); break; default: return 0; @@ -4119,7 +4153,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;