From 858a156cfb11d70aba0a6fc28a15906fdceaa429 Mon Sep 17 00:00:00 2001
From: Carl Meyer <carl@oddbird.net>
Date: Tue, 29 Aug 2023 17:40:35 -0700
Subject: [PATCH] gh-108654: restore comprehension locals before handling
 exception

---
 Lib/test/test_listcomps.py | 25 ++++++++++++++
 Python/compile.c           | 69 +++++++++++++++++++++++++++++---------
 2 files changed, 79 insertions(+), 15 deletions(-)

diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index 9f28ced32bd26c6..9af7a54302547d6 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/Python/compile.c b/Python/compile.c
index 6b816b4c6eda6c0..502a43c6ed46a6a 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);