From 627ff93c192a6f5e7151903699a5be17dd15ed63 Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Tue, 29 Aug 2023 17:40:35 -0700 Subject: [PATCH 1/5] gh-108654: restore comprehension locals before handling exception --- Lib/test/test_listcomps.py | 25 +++++++ ...-08-29-17-53-12.gh-issue-108654.jbkDVo.rst | 2 + Python/compile.c | 69 +++++++++++++++---- 3 files changed, 81 insertions(+), 15 deletions(-) create mode 100644 Misc/NEWS.d/next/Core and Builtins/2023-08-29-17-53-12.gh-issue-108654.jbkDVo.rst diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py index 9f28ced32bd26c..9af7a54302547d 100644 --- a/Lib/test/test_listcomps.py +++ b/Lib/test/test_listcomps.py @@ -561,6 +561,31 @@ def test_iter_var_available_in_locals(self): } ) + 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"]}) + __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 00000000000000..032e0331b20e75 --- /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 6b816b4c6eda6c..502a43c6ed46a6 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 @@ -5641,6 +5643,42 @@ push_inlined_comprehension_state(struct compiler *c, location loc, ADDOP_I(c, loc, SWAP, PyList_GET_SIZE(state->pushed_locals) + 1); } + if (c->u->u_nfblocks > 0) { + // If we are inside a try block, we need to add our own cleanup handler + // to restore comprehension locals, 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; } @@ -5651,6 +5689,20 @@ pop_inlined_comprehension_state(struct compiler *c, location loc, c->u->u_in_inlined_comp--; PyObject *k, *v; Py_ssize_t pos = 0; + if (IS_LABEL(state.cleanup)) { + 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); + restore_inlined_comprehension_locals(c, loc, state); + ADDOP_I(c, NO_LOCATION, RERAISE, 0); + ADDOP(c, NO_LOCATION, POP_BLOCK); + + USE_LABEL(c, state.end); + } if (state.temp_symbols) { while (PyDict_Next(state.temp_symbols, &pos, &k, &v)) { if (PyDict_SetItem(c->u->u_ste->ste_symbols, k, v)) { @@ -5660,20 +5712,7 @@ 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); - } + restore_inlined_comprehension_locals(c, loc, state); Py_CLEAR(state.pushed_locals); } if (state.fast_hidden) { @@ -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); From c8eb6bff5166e81f2b99256d199985e10b157848 Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Wed, 30 Aug 2023 06:08:56 -0600 Subject: [PATCH 2/5] Apply suggestions from code review Co-authored-by: Dong-hee Na --- Python/compile.c | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/Python/compile.c b/Python/compile.c index 502a43c6ed46a6..375d17e886c780 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -5697,7 +5697,9 @@ pop_inlined_comprehension_state(struct compiler *c, location loc, // discard incomplete comprehension result (beneath exc on stack) ADDOP_I(c, NO_LOCATION, SWAP, 2); ADDOP(c, NO_LOCATION, POP_TOP); - restore_inlined_comprehension_locals(c, loc, state); + if (restore_inlined_comprehension_locals(c, loc, state) < 0) { + return ERROR; + } ADDOP_I(c, NO_LOCATION, RERAISE, 0); ADDOP(c, NO_LOCATION, POP_BLOCK); @@ -5712,7 +5714,9 @@ pop_inlined_comprehension_state(struct compiler *c, location loc, Py_CLEAR(state.temp_symbols); } if (state.pushed_locals) { - restore_inlined_comprehension_locals(c, loc, state); + if (restore_inlined_comprehension_locals(c, loc, state) < 0) { + return ERROR; + } Py_CLEAR(state.pushed_locals); } if (state.fast_hidden) { From 4c4eebc7bac0c74379971eca55942dd722434ff4 Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Wed, 30 Aug 2023 05:25:55 -0700 Subject: [PATCH 3/5] fix location of POP_BLOCK --- Lib/test/test_listcomps.py | 10 ++++++++++ Python/compile.c | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py index 9af7a54302547d..bedd99b4a44fcb 100644 --- a/Lib/test/test_listcomps.py +++ b/Lib/test/test_listcomps.py @@ -586,6 +586,16 @@ def f(value): """ 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/Python/compile.c b/Python/compile.c index 375d17e886c780..8722164c700717 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -5690,6 +5690,7 @@ pop_inlined_comprehension_state(struct compiler *c, location loc, PyObject *k, *v; Py_ssize_t pos = 0; if (IS_LABEL(state.cleanup)) { + ADDOP(c, NO_LOCATION, POP_BLOCK); ADDOP_JUMP(c, NO_LOCATION, JUMP, state.end); // cleanup from an exception inside the comprehension @@ -5701,7 +5702,6 @@ pop_inlined_comprehension_state(struct compiler *c, location loc, return ERROR; } ADDOP_I(c, NO_LOCATION, RERAISE, 0); - ADDOP(c, NO_LOCATION, POP_BLOCK); USE_LABEL(c, state.end); } From 8e8b41f3ddedac4465f617bde52bbf8bfff5d081 Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Wed, 30 Aug 2023 05:51:11 -0700 Subject: [PATCH 4/5] no cleanup needed if no pushed locals --- Python/compile.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Python/compile.c b/Python/compile.c index 8722164c700717..b9b7f387fc6156 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -5643,7 +5643,7 @@ push_inlined_comprehension_state(struct compiler *c, location loc, ADDOP_I(c, loc, SWAP, PyList_GET_SIZE(state->pushed_locals) + 1); } - if (c->u->u_nfblocks > 0) { + if (c->u->u_nfblocks > 0 && state->pushed_locals) { // If we are inside a try block, we need to add our own cleanup handler // to restore comprehension locals, so they have the correct values // inside an exception handler or finally block. From d49e911fc5243695c17acd0ec0602a4126f44727 Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Wed, 30 Aug 2023 10:35:53 -0700 Subject: [PATCH 5/5] always restore comp-shadowed locals on exception --- Python/compile.c | 28 ++++++++++++---------------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/Python/compile.c b/Python/compile.c index b9b7f387fc6156..50e29b4607f7ce 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -5641,12 +5641,10 @@ 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); - } - if (c->u->u_nfblocks > 0 && state->pushed_locals) { - // If we are inside a try block, we need to add our own cleanup handler - // to restore comprehension locals, so they have the correct values - // inside an exception handler or finally block. + // 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); @@ -5689,7 +5687,15 @@ pop_inlined_comprehension_state(struct compiler *c, location loc, c->u->u_in_inlined_comp--; PyObject *k, *v; Py_ssize_t pos = 0; - if (IS_LABEL(state.cleanup)) { + if (state.temp_symbols) { + while (PyDict_Next(state.temp_symbols, &pos, &k, &v)) { + if (PyDict_SetItem(c->u->u_ste->ste_symbols, k, v)) { + return ERROR; + } + } + Py_CLEAR(state.temp_symbols); + } + if (state.pushed_locals) { ADDOP(c, NO_LOCATION, POP_BLOCK); ADDOP_JUMP(c, NO_LOCATION, JUMP, state.end); @@ -5704,16 +5710,6 @@ pop_inlined_comprehension_state(struct compiler *c, location loc, ADDOP_I(c, NO_LOCATION, RERAISE, 0); USE_LABEL(c, state.end); - } - if (state.temp_symbols) { - while (PyDict_Next(state.temp_symbols, &pos, &k, &v)) { - if (PyDict_SetItem(c->u->u_ste->ste_symbols, k, v)) { - return ERROR; - } - } - Py_CLEAR(state.temp_symbols); - } - if (state.pushed_locals) { if (restore_inlined_comprehension_locals(c, loc, state) < 0) { return ERROR; }