diff --git a/py/compile.c b/py/compile.c index 5c2fce686600352bdda210b6281e893166c6eac6..841b8f90c0fc833626beacb0bff5948c094d03b8 100644 --- a/py/compile.c +++ b/py/compile.c @@ -1421,6 +1421,19 @@ STATIC void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) { optimize = false; } } + // arguments must be able to be compiled as standard expressions + if (optimize && MP_PARSE_NODE_IS_STRUCT(pn_range_start)) { + int k = MP_PARSE_NODE_STRUCT_KIND((mp_parse_node_struct_t*)pn_range_start); + if (k == PN_arglist_star || k == PN_arglist_dbl_star || k == PN_argument) { + optimize = false; + } + } + if (optimize && MP_PARSE_NODE_IS_STRUCT(pn_range_end)) { + int k = MP_PARSE_NODE_STRUCT_KIND((mp_parse_node_struct_t*)pn_range_end); + if (k == PN_arglist_star || k == PN_arglist_dbl_star || k == PN_argument) { + optimize = false; + } + } } if (optimize) { compile_for_stmt_optimised_range(comp, pns->nodes[0], pn_range_start, pn_range_end, pn_range_step, pns->nodes[2], pns->nodes[3]); diff --git a/tests/basics/for_range.py b/tests/basics/for_range.py new file mode 100644 index 0000000000000000000000000000000000000000..ddff5ebd44357e5c4d25610e188ec6de85e94996 --- /dev/null +++ b/tests/basics/for_range.py @@ -0,0 +1,58 @@ +# test for+range, mostly to check optimisation of this pair + +# apply args using * +for x in range(*(1, 3)): + print(x) +for x in range(1, *(6, 2)): + print(x) + +# apply args using ** +try: + for x in range(**{'end':1}): + print(x) +except TypeError: + print('TypeError') +try: + for x in range(0, **{'end':1}): + print(x) +except TypeError: + print('TypeError') +try: + for x in range(0, 1, **{'step':1}): + print(x) +except TypeError: + print('TypeError') + +# keyword args +try: + for x in range(end=1): + print(x) +except TypeError: + print('TypeError') +try: + for x in range(0, end=1): + print(x) +except TypeError: + print('TypeError') +try: + for x in range(0, 1, step=1): + print(x) +except TypeError: + print('TypeError') + +# argument is a comprehension +try: + for x in range(0 for i in []): + print(x) +except TypeError: + print('TypeError') +try: + for x in range(0, (0 for i in [])): + print(x) +except TypeError: + print('TypeError') +try: + for x in range(0, 1, (0 for i in [])): + print(x) +except TypeError: + print('TypeError')