diff options
| -rw-r--r-- | Lib/test/test_listcomps.py | 35 | ||||
| -rw-r--r-- | Misc/NEWS.d/next/Core and Builtins/2023-08-29-17-53-12.gh-issue-108654.jbkDVo.rst | 2 | ||||
| -rw-r--r-- | Python/compile.c | 67 |
3 files changed, 90 insertions, 14 deletions
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py index 9f28ced..bedd99b 100644 --- a/Lib/test/test_listcomps.py +++ b/Lib/test/test_listcomps.py @@ -561,6 +561,41 @@ class ListComprehensionTest(unittest.TestCase): } ) + def test_comp_in_try_except(self): + template = """ + value = ["a"] + try: + [{func}(value) for value in value] + except: + pass + """ + for func in ["str", "int"]: + code = template.format(func=func) + raises = func != "str" + with self.subTest(raises=raises): + self._check_in_scopes(code, {"value": ["a"]}) + + def test_comp_in_try_finally(self): + code = """ + def f(value): + try: + [{func}(value) for value in value] + finally: + return value + ret = f(["a"]) + """ + self._check_in_scopes(code, {"ret": ["a"]}) + + def test_exception_in_post_comp_call(self): + code = """ + value = [1, None] + try: + [v for v in value].sort() + except: + pass + """ + self._check_in_scopes(code, {"value": [1, None]}) + __test__ = {'doctests' : doctests} diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-08-29-17-53-12.gh-issue-108654.jbkDVo.rst b/Misc/NEWS.d/next/Core and Builtins/2023-08-29-17-53-12.gh-issue-108654.jbkDVo.rst new file mode 100644 index 0000000..032e033 --- /dev/null +++ b/Misc/NEWS.d/next/Core and Builtins/2023-08-29-17-53-12.gh-issue-108654.jbkDVo.rst @@ -0,0 +1,2 @@ +Restore locals shadowed by an inlined comprehension if the comprehension +raises an exception. diff --git a/Python/compile.c b/Python/compile.c index 6b816b4..50e29b4 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -5529,6 +5529,8 @@ typedef struct { PyObject *pushed_locals; PyObject *temp_symbols; PyObject *fast_hidden; + jump_target_label cleanup; + jump_target_label end; } inlined_comprehension_state; static int @@ -5639,12 +5641,46 @@ push_inlined_comprehension_state(struct compiler *c, location loc, // `pushed_locals` on the stack, but this will be reversed when we swap // out the comprehension result in pop_inlined_comprehension_state ADDOP_I(c, loc, SWAP, PyList_GET_SIZE(state->pushed_locals) + 1); + + // Add our own cleanup handler to restore comprehension locals in case + // of exception, so they have the correct values inside an exception + // handler or finally block. + NEW_JUMP_TARGET_LABEL(c, cleanup); + state->cleanup = cleanup; + NEW_JUMP_TARGET_LABEL(c, end); + state->end = end; + + // no need to push an fblock for this "virtual" try/finally; there can't + // be return/continue/break inside a comprehension + ADDOP_JUMP(c, loc, SETUP_FINALLY, cleanup); } return SUCCESS; } static int +restore_inlined_comprehension_locals(struct compiler *c, location loc, + inlined_comprehension_state state) +{ + PyObject *k; + // pop names we pushed to stack earlier + Py_ssize_t npops = PyList_GET_SIZE(state.pushed_locals); + // Preserve the comprehension result (or exception) as TOS. This + // reverses the SWAP we did in push_inlined_comprehension_state to get + // the outermost iterable to TOS, so we can still just iterate + // pushed_locals in simple reverse order + ADDOP_I(c, loc, SWAP, npops + 1); + for (Py_ssize_t i = npops - 1; i >= 0; --i) { + k = PyList_GetItem(state.pushed_locals, i); + if (k == NULL) { + return ERROR; + } + ADDOP_NAME(c, loc, STORE_FAST_MAYBE_NULL, k, varnames); + } + return SUCCESS; +} + +static int pop_inlined_comprehension_state(struct compiler *c, location loc, inlined_comprehension_state state) { @@ -5660,19 +5696,22 @@ pop_inlined_comprehension_state(struct compiler *c, location loc, Py_CLEAR(state.temp_symbols); } if (state.pushed_locals) { - // pop names we pushed to stack earlier - Py_ssize_t npops = PyList_GET_SIZE(state.pushed_locals); - // Preserve the list/dict/set result of the comprehension as TOS. This - // reverses the SWAP we did in push_inlined_comprehension_state to get - // the outermost iterable to TOS, so we can still just iterate - // pushed_locals in simple reverse order - ADDOP_I(c, loc, SWAP, npops + 1); - for (Py_ssize_t i = npops - 1; i >= 0; --i) { - k = PyList_GetItem(state.pushed_locals, i); - if (k == NULL) { - return ERROR; - } - ADDOP_NAME(c, loc, STORE_FAST_MAYBE_NULL, k, varnames); + ADDOP(c, NO_LOCATION, POP_BLOCK); + ADDOP_JUMP(c, NO_LOCATION, JUMP, state.end); + + // cleanup from an exception inside the comprehension + USE_LABEL(c, state.cleanup); + // discard incomplete comprehension result (beneath exc on stack) + ADDOP_I(c, NO_LOCATION, SWAP, 2); + ADDOP(c, NO_LOCATION, POP_TOP); + if (restore_inlined_comprehension_locals(c, loc, state) < 0) { + return ERROR; + } + ADDOP_I(c, NO_LOCATION, RERAISE, 0); + + USE_LABEL(c, state.end); + if (restore_inlined_comprehension_locals(c, loc, state) < 0) { + return ERROR; } Py_CLEAR(state.pushed_locals); } @@ -5715,7 +5754,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type, expr_ty val) { PyCodeObject *co = NULL; - inlined_comprehension_state inline_state = {NULL, NULL}; + inlined_comprehension_state inline_state = {NULL, NULL, NULL, NO_LABEL, NO_LABEL}; comprehension_ty outermost; int scope_type = c->u->u_scope_type; int is_top_level_await = IS_TOP_LEVEL_AWAIT(c); |