gh-104619: never leak comprehension locals to outer locals() (#104637)

python · May 19, 2023 · 70c7796 · 70c7796
1 parent 86e6f16
commit 70c7796
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 21 deletions.
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
@@ -516,6 +516,19 @@ def test_assign_to_comp_iter_var_in_outer_function(self):
         """
         self._check_in_scopes(code, {"a": [1]}, scopes=["function"])
 
+    def test_no_leakage_to_locals(self):
+        code = """
+            def b():
+                [a for b in [1] for _ in []]
+                return b, locals()
+            r, s = b()
+            x = r is b
+            y = list(s.keys())
+        """
+        self._check_in_scopes(code, {"x": True, "y": []}, scopes=["module"])
+        self._check_in_scopes(code, {"x": True, "y": ["b"]}, scopes=["function"])
+        self._check_in_scopes(code, raises=NameError, scopes=["class"])
+
 
 __test__ = {'doctests' : doctests}
 

diff --git a/Python/compile.c b/Python/compile.c
@@ -5490,31 +5490,29 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
                 }
                 Py_DECREF(outv);
             }
-            if (outsc == LOCAL || outsc == CELL || outsc == FREE) {
-                // local names bound in comprehension must be isolated from
-                // outer scope; push existing value (which may be NULL if
-                // not defined) on stack
+            // local names bound in comprehension must be isolated from
+            // outer scope; push existing value (which may be NULL if
+            // not defined) on stack
+            if (state->pushed_locals == NULL) {
+                state->pushed_locals = PyList_New(0);
                 if (state->pushed_locals == NULL) {
-                    state->pushed_locals = PyList_New(0);
-                    if (state->pushed_locals == NULL) {
-                        return ERROR;
-                    }
-                }
-                // in the case of a cell, this will actually push the cell
-                // itself to the stack, then we'll create a new one for the
-                // comprehension and restore the original one after
-                ADDOP_NAME(c, loc, LOAD_FAST_AND_CLEAR, k, varnames);
-                if (scope == CELL) {
-                    if (outsc == FREE) {
-                        ADDOP_NAME(c, loc, MAKE_CELL, k, freevars);
-                    } else {
-                        ADDOP_NAME(c, loc, MAKE_CELL, k, cellvars);
-                    }
-                }
-                if (PyList_Append(state->pushed_locals, k) < 0) {
                     return ERROR;
                 }
             }
+            // in the case of a cell, this will actually push the cell
+            // itself to the stack, then we'll create a new one for the
+            // comprehension and restore the original one after
+            ADDOP_NAME(c, loc, LOAD_FAST_AND_CLEAR, k, varnames);
+            if (scope == CELL) {
+                if (outsc == FREE) {
+                    ADDOP_NAME(c, loc, MAKE_CELL, k, freevars);
+                } else {
+                    ADDOP_NAME(c, loc, MAKE_CELL, k, cellvars);
+                }
+            }
+            if (PyList_Append(state->pushed_locals, k) < 0) {
+                return ERROR;
+            }
         }
     }
     if (state->pushed_locals) {