gh-105340: include hidden fast-locals in locals()

Author: carljm

Include/cpython/ceval.h

@@ -31,3 +31,5 @@ _PyEval_RequestCodeExtraIndex(freefunc f) {
 
 PyAPI_FUNC(int) _PyEval_SliceIndex(PyObject *, Py_ssize_t *);
 PyAPI_FUNC(int) _PyEval_SliceIndexNotNone(PyObject *, Py_ssize_t *);
+
+PyAPI_FUNC(PyObject *) PyEval_GetFrameLocals(void);

Include/internal/pycore_frame.h

@@ -226,6 +226,9 @@ _PyFrame_ClearExceptCode(_PyInterpreterFrame * frame);
 int
 _PyFrame_Traverse(_PyInterpreterFrame *frame, visitproc visit, void *arg);
 
+PyObject *
+_PyFrame_GetLocals(_PyInterpreterFrame *frame, int include_hidden);
+
 int
 _PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame);
 

Lib/test/test_listcomps.py

@@ -539,6 +539,19 @@ def b():
         self._check_in_scopes(code, {"x": True, "y": ["b"]}, scopes=["function"])
         self._check_in_scopes(code, raises=NameError, scopes=["class"])
 
+    def test_iter_var_available_in_locals(self):
+        code = """
+            l = [1, 2]
+            y = 0
+            items = [locals()["x"] for x in l]
+            items2 = [vars()["x"] for x in l]
+            items3 = [eval("x") for x in l]
+            # x is available, and does not overwrite y
+            [exec("y = x") for x in l]
+        """
+        self._check_in_scopes(
+            code, {"items": [1, 2], "items2": [1, 2], "items3": [1, 2], "y": 0})
+
 
 __test__ = {'doctests' : doctests}
 

Misc/NEWS.d/next/Core and Builtins/2023-06-12-16-38-31.gh-issue-105340._jRHXe.rst

@@ -0,0 +1,2 @@
+Include the comprehension iteration variable in ``locals()`` inside a
+module- or class-scope comprehension.

Objects/frameobject.c

@@ -1199,15 +1199,28 @@ frame_get_var(_PyInterpreterFrame *frame, PyCodeObject *co, int i,
     return 1;
 }
 
-int
-_PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame)
+
+PyObject *
+_PyFrame_GetLocals(_PyInterpreterFrame *frame, int include_hidden)
 {
     /* Merge fast locals into f->f_locals */
     PyObject *locals = frame->f_locals;
     if (locals == NULL) {
         locals = frame->f_locals = PyDict_New();
         if (locals == NULL) {
-            return -1;
+            return NULL;
+        }
+    }
+    PyObject *hidden = NULL;
+
+    /* If include_hidden, "hidden" fast locals (from inlined comprehensions in
+       module/class scopes) will be included in the returned dict, but not in
+       frame->f_locals; the returned dict will be a modified copy. Non-hidden
+       locals will still be updated in frame->f_locals. */
+    if (include_hidden) {
+        hidden = PyDict_New();
+        if (hidden == NULL) {
+            return NULL;
         }
     }
 
@@ -1223,6 +1236,11 @@ _PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame)
         PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
         _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
         if (kind & CO_FAST_HIDDEN) {
+            if (include_hidden && value != NULL) {
+                if (PyObject_SetItem(hidden, name, value) != 0) {
+                    goto error;
+                }
+            }
             continue;
         }
         if (value == NULL) {
@@ -1231,16 +1249,53 @@ _PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame)
                     PyErr_Clear();
                 }
                 else {
-                    return -1;
+                    goto error;
                 }
             }
         }
         else {
             if (PyObject_SetItem(locals, name, value) != 0) {
-                return -1;
+                goto error;
             }
         }
     }
+
+    if (include_hidden && PyDict_Size(hidden)) {
+        PyObject *innerlocals = PyDict_New();
+        if (innerlocals == NULL) {
+            goto error;
+        }
+        if (PyDict_Merge(innerlocals, locals, 1) != 0) {
+            Py_DECREF(innerlocals);
+            goto error;
+        }
+        if (PyDict_Merge(innerlocals, hidden, 1) != 0) {
+            Py_DECREF(innerlocals);
+            goto error;
+        }
+        locals = innerlocals;
+    }
+    else {
+        Py_INCREF(locals);
+    }
+    Py_CLEAR(hidden);
+
+    return locals;
+
+  error:
+    Py_XDECREF(hidden);
+    return NULL;
+}
+
+
+int
+_PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame)
+{
+    PyObject *locals = _PyFrame_GetLocals(frame, 0);
+    if (locals == NULL) {
+        return -1;
+    }
+    Py_DECREF(locals);
     return 0;
 }
 

Python/bltinmodule.c

@@ -907,7 +907,7 @@ builtin_eval_impl(PyObject *module, PyObject *source, PyObject *globals,
                   PyObject *locals)
 /*[clinic end generated code: output=0a0824aa70093116 input=11ee718a8640e527]*/
 {
-    PyObject *result, *source_copy;
+    PyObject *result = NULL, *source_copy;
     const char *str;
 
     if (locals != Py_None && !PyMapping_Check(locals)) {
@@ -923,54 +923,61 @@ builtin_eval_impl(PyObject *module, PyObject *source, PyObject *globals,
     if (globals == Py_None) {
         globals = PyEval_GetGlobals();
         if (locals == Py_None) {
-            locals = PyEval_GetLocals();
+            locals = PyEval_GetFrameLocals();
             if (locals == NULL)
                 return NULL;
         }
     }
     else if (locals == Py_None)
-        locals = globals;
+        locals = Py_NewRef(globals);
+    else {
+        Py_INCREF(locals);
+    }
 
     if (globals == NULL || locals == NULL) {
         PyErr_SetString(PyExc_TypeError,
             "eval must be given globals and locals "
             "when called without a frame");
-        return NULL;
+        goto error;
     }
 
     int r = PyDict_Contains(globals, &_Py_ID(__builtins__));
     if (r == 0) {
         r = PyDict_SetItem(globals, &_Py_ID(__builtins__), PyEval_GetBuiltins());
     }
     if (r < 0) {
-        return NULL;
+        goto error;
     }
 
     if (PyCode_Check(source)) {
         if (PySys_Audit("exec", "O", source) < 0) {
-            return NULL;
+            goto error;
         }
 
         if (PyCode_GetNumFree((PyCodeObject *)source) > 0) {
             PyErr_SetString(PyExc_TypeError,
                 "code object passed to eval() may not contain free variables");
-            return NULL;
+            goto error;
         }
-        return PyEval_EvalCode(source, globals, locals);
+        result = PyEval_EvalCode(source, globals, locals);
     }
+    else {
+        PyCompilerFlags cf = _PyCompilerFlags_INIT;
+        cf.cf_flags = PyCF_SOURCE_IS_UTF8;
+        str = _Py_SourceAsString(source, "eval", "string, bytes or code", &cf, &source_copy);
+        if (str == NULL)
+            goto error;
 
-    PyCompilerFlags cf = _PyCompilerFlags_INIT;
-    cf.cf_flags = PyCF_SOURCE_IS_UTF8;
-    str = _Py_SourceAsString(source, "eval", "string, bytes or code", &cf, &source_copy);
-    if (str == NULL)
-        return NULL;
+        while (*str == ' ' || *str == '\t')
+            str++;
 
-    while (*str == ' ' || *str == '\t')
-        str++;
+        (void)PyEval_MergeCompilerFlags(&cf);
+        result = PyRun_StringFlags(str, Py_eval_input, globals, locals, &cf);
+        Py_XDECREF(source_copy);
+    }
 
-    (void)PyEval_MergeCompilerFlags(&cf);
-    result = PyRun_StringFlags(str, Py_eval_input, globals, locals, &cf);
-    Py_XDECREF(source_copy);
+  error:
+    Py_XDECREF(locals);
     return result;
 }
 
@@ -1005,7 +1012,7 @@ builtin_exec_impl(PyObject *module, PyObject *source, PyObject *globals,
     if (globals == Py_None) {
         globals = PyEval_GetGlobals();
         if (locals == Py_None) {
-            locals = PyEval_GetLocals();
+            locals = PyEval_GetFrameLocals();
             if (locals == NULL)
                 return NULL;
         }
@@ -1015,26 +1022,30 @@ builtin_exec_impl(PyObject *module, PyObject *source, PyObject *globals,
             return NULL;
         }
     }
-    else if (locals == Py_None)
-        locals = globals;
+    else if (locals == Py_None) {
+        locals = Py_NewRef(globals);
+    }
+    else {
+        Py_INCREF(locals);
+    }
 
     if (!PyDict_Check(globals)) {
         PyErr_Format(PyExc_TypeError, "exec() globals must be a dict, not %.100s",
                      Py_TYPE(globals)->tp_name);
-        return NULL;
+        goto error;
     }
     if (!PyMapping_Check(locals)) {
         PyErr_Format(PyExc_TypeError,
             "locals must be a mapping or None, not %.100s",
             Py_TYPE(locals)->tp_name);
-        return NULL;
+        goto error;
     }
     int r = PyDict_Contains(globals, &_Py_ID(__builtins__));
     if (r == 0) {
         r = PyDict_SetItem(globals, &_Py_ID(__builtins__), PyEval_GetBuiltins());
     }
     if (r < 0) {
-        return NULL;
+        goto error;
     }
 
     if (closure == Py_None) {
@@ -1047,7 +1058,7 @@ builtin_exec_impl(PyObject *module, PyObject *source, PyObject *globals,
             if (closure) {
                 PyErr_SetString(PyExc_TypeError,
                     "cannot use a closure with this code object");
-                return NULL;
+                goto error;
             }
         } else {
             int closure_is_ok =
@@ -1067,12 +1078,12 @@ builtin_exec_impl(PyObject *module, PyObject *source, PyObject *globals,
                 PyErr_Format(PyExc_TypeError,
                     "code object requires a closure of exactly length %zd",
                     num_free);
-                return NULL;
+                goto error;
             }
         }
 
         if (PySys_Audit("exec", "O", source) < 0) {
-            return NULL;
+            goto error;
         }
 
         if (!closure) {
@@ -1099,7 +1110,7 @@ builtin_exec_impl(PyObject *module, PyObject *source, PyObject *globals,
                                        "string, bytes or code", &cf,
                                        &source_copy);
         if (str == NULL)
-            return NULL;
+            goto error;
         if (PyEval_MergeCompilerFlags(&cf))
             v = PyRun_StringFlags(str, Py_file_input, globals,
                                   locals, &cf);
@@ -1108,9 +1119,14 @@ builtin_exec_impl(PyObject *module, PyObject *source, PyObject *globals,
         Py_XDECREF(source_copy);
     }
     if (v == NULL)
-        return NULL;
+        goto error;
+    Py_DECREF(locals);
     Py_DECREF(v);
     Py_RETURN_NONE;
+
+  error:
+    Py_XDECREF(locals);
+    return NULL;
 }
 
 
@@ -1722,8 +1738,7 @@ builtin_locals_impl(PyObject *module)
 {
     PyObject *d;
 
-    d = PyEval_GetLocals();
-    return Py_XNewRef(d);
+    return PyEval_GetFrameLocals();
 }
 
 
@@ -2441,7 +2456,7 @@ builtin_vars_impl(PyObject *module, PyObject *object)
     PyObject *d;
 
     if (object == NULL) {
-        d = Py_XNewRef(PyEval_GetLocals());
+        d = PyEval_GetFrameLocals();
     }
     else {
         if (_PyObject_LookupAttr(object, &_Py_ID(__dict__), &d) == 0) {

Python/ceval.c

@@ -2257,6 +2257,19 @@ PyEval_GetLocals(void)
     return locals;
 }
 
+PyObject *
+PyEval_GetFrameLocals(void)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+     _PyInterpreterFrame *current_frame = _PyThreadState_GetFrame(tstate);
+    if (current_frame == NULL) {
+        _PyErr_SetString(tstate, PyExc_SystemError, "frame does not exist");
+        return NULL;
+    }
+
+    return _PyFrame_GetLocals(current_frame, 1);
+}
+
 PyObject *
 PyEval_GetGlobals(void)
 {