gh-91248: Optimize PyFrame_GetVar() (#99252)

PyFrame_GetVar() no longer creates a temporary dictionary to get a
variable.

Doc/c-api/frame.rst

@@ -87,6 +87,8 @@ See also :ref:`Reflection <reflection>`.
    * Raise :exc:`NameError` and return ``NULL`` if the variable does not exist.
    * Raise an exception and return ``NULL`` on error.
 
+   *name* type must be a :class:`str`.
+
    .. versionadded:: 3.12
 
 

Lib/test/test_frame.py

@@ -352,6 +352,12 @@ class TestCAPI(unittest.TestCase):
         with self.assertRaises(NameError):
             _testcapi.frame_getvarstring(current_frame, b"y")
 
+        # wrong name type
+        with self.assertRaises(TypeError):
+            _testcapi.frame_getvar(current_frame, b'x')
+        with self.assertRaises(TypeError):
+            _testcapi.frame_getvar(current_frame, 123)
+
     def getgenframe(self):
         yield sys._getframe()
 

Objects/frameobject.c

@@ -1109,28 +1109,23 @@ _PyFrame_OpAlreadyRan(_PyInterpreterFrame *frame, int opcode, int oparg)
     return 0;
 }
 
-int
+
-_PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame) {
+// Initialize frame free variables if needed
-    /* Merge fast locals into f->f_locals */
+static void
-    PyObject *locals;
+frame_init_get_vars(_PyInterpreterFrame *frame)
-    PyObject **fast;
+{
-    PyCodeObject *co;
-    locals = frame->f_locals;
-    if (locals == NULL) {
-        locals = frame->f_locals = PyDict_New();
-        if (locals == NULL)
-            return -1;
-    }
-    co = frame->f_code;
-    fast = _PyFrame_GetLocalsArray(frame);
     // COPY_FREE_VARS has no quickened forms, so no need to use _PyOpcode_Deopt
     // here:
-    int lasti = _PyInterpreterFrame_LASTI(frame);
-    if (lasti < 0 && _Py_OPCODE(_PyCode_CODE(co)[0]) == COPY_FREE_VARS
-        && PyFunction_Check(frame->f_funcobj))
-    {
-        /* Free vars have not been initialized -- Do that */
     PyCodeObject *co = frame->f_code;
+    int lasti = _PyInterpreterFrame_LASTI(frame);
+    if (!(lasti < 0 && _Py_OPCODE(_PyCode_CODE(co)[0]) == COPY_FREE_VARS
+          && PyFunction_Check(frame->f_funcobj)))
+    {
+        /* Free vars are initialized */
+        return;
+    }
+
+    /* Free vars have not been initialized -- Do that */
     PyObject *closure = ((PyFunctionObject *)frame->f_funcobj)->func_closure;
     int offset = co->co_nlocals + co->co_nplaincellvars;
     for (int i = 0; i < co->co_nfreevars; ++i) {
@@ -1139,8 +1134,13 @@ _PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame) {
     }
     // COPY_FREE_VARS doesn't have inline CACHEs, either:
     frame->prev_instr = _PyCode_CODE(frame->f_code);
-    }
+}
-    for (int i = 0; i < co->co_nlocalsplus; i++) {
+
+
+static int
+frame_get_var(_PyInterpreterFrame *frame, PyCodeObject *co, int i,
+              PyObject **pvalue)
+{
     _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
 
     /* If the namespace is unoptimized, then one of the
@@ -1152,11 +1152,10 @@ _PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame) {
        into the locals dict used by the class.
     */
     if (kind & CO_FAST_FREE && !(co->co_flags & CO_OPTIMIZED)) {
-            continue;
+        return 0;
     }
 
-        PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
+    PyObject *value = frame->localsplus[i];
-        PyObject *value = fast[i];
     if (frame->stacktop) {
         if (kind & CO_FAST_FREE) {
             // The cell was set by COPY_FREE_VARS.
@@ -1184,6 +1183,32 @@ _PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame) {
     else {
         assert(value == NULL);
     }
+    *pvalue = value;
+    return 1;
+}
+
+int
+_PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame)
+{
+    /* 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;
+        }
+    }
+
+    frame_init_get_vars(frame);
+
+    PyCodeObject *co = frame->f_code;
+    for (int i = 0; i < co->co_nlocalsplus; i++) {
+        PyObject *value;  // borrowed reference
+        if (!frame_get_var(frame, co, i, &value)) {
+            continue;
+        }
+
+        PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
         if (value == NULL) {
             if (PyObject_DelItem(locals, name) != 0) {
                 if (PyErr_ExceptionMatches(PyExc_KeyError)) {
@@ -1203,6 +1228,54 @@ _PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame) {
     return 0;
 }
 
+
+PyObject *
+PyFrame_GetVar(PyFrameObject *frame_obj, PyObject *name)
+{
+    if (!PyUnicode_Check(name)) {
+        PyErr_Format(PyExc_TypeError, "name must be str, not %s",
+                     Py_TYPE(name)->tp_name);
+        return NULL;
+    }
+
+    _PyInterpreterFrame *frame = frame_obj->f_frame;
+    frame_init_get_vars(frame);
+
+    PyCodeObject *co = frame->f_code;
+    for (int i = 0; i < co->co_nlocalsplus; i++) {
+        PyObject *var_name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
+        if (!_PyUnicode_Equal(var_name, name)) {
+            continue;
+        }
+
+        PyObject *value;  // borrowed reference
+        if (!frame_get_var(frame, co, i, &value)) {
+            break;
+        }
+        if (value == NULL) {
+            break;
+        }
+        return Py_NewRef(value);
+    }
+
+    PyErr_Format(PyExc_NameError, "variable %R does not exist", name);
+    return NULL;
+}
+
+
+PyObject *
+PyFrame_GetVarString(PyFrameObject *frame, const char *name)
+{
+    PyObject *name_obj = PyUnicode_FromString(name);
+    if (name_obj == NULL) {
+        return NULL;
+    }
+    PyObject *value = PyFrame_GetVar(frame, name_obj);
+    Py_DECREF(name_obj);
+    return value;
+}
+
+
 int
 PyFrame_FastToLocalsWithError(PyFrameObject *f)
 {
@@ -1413,35 +1486,3 @@ _PyEval_BuiltinsFromGlobals(PyThreadState *tstate, PyObject *globals)
 
     return _PyEval_GetBuiltins(tstate);
 }
-
-PyObject *
-PyFrame_GetVar(PyFrameObject *frame, PyObject *name)
-{
-    PyObject *locals = PyFrame_GetLocals(frame);
-    if (locals == NULL) {
-        return NULL;
-    }
-    PyObject *value = PyDict_GetItemWithError(locals, name);
-    Py_DECREF(locals);
-
-    if (value == NULL) {
-        if (PyErr_Occurred()) {
-            return NULL;
-        }
-        PyErr_Format(PyExc_NameError, "variable %R does not exist", name);
-        return NULL;
-    }
-    return Py_NewRef(value);
-}
-
-PyObject *
-PyFrame_GetVarString(PyFrameObject *frame, const char *name)
-{
-    PyObject *name_obj = PyUnicode_FromString(name);
-    if (name_obj == NULL) {
-        return NULL;
-    }
-    PyObject *value = PyFrame_GetVar(frame, name_obj);
-    Py_DECREF(name_obj);
-    return value;
-}