pythongh-125590: Allow FrameLocalsProxy to delete and pop keys from extra locals (pythonGH-125616)

(cherry picked from commit 5b7a872)

Co-authored-by: Tian Gao <[email protected]>

Author: gaogaotiantian

Lib/test/test_frame.py

@@ -399,15 +399,41 @@ def test_repr(self):
     def test_delete(self):
         x = 1
         d = sys._getframe().f_locals
-        with self.assertRaises(TypeError):
+
+        # This needs to be tested before f_extra_locals is created
+        with self.assertRaisesRegex(KeyError, 'non_exist'):
+            del d['non_exist']
+
+        with self.assertRaises(KeyError):
+            d.pop('non_exist')
+
+        with self.assertRaisesRegex(ValueError, 'local variables'):
             del d['x']
 
         with self.assertRaises(AttributeError):
             d.clear()
 
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(ValueError):
             d.pop('x')
 
+        with self.assertRaises(ValueError):
+            d.pop('x', None)
+
+        # 'm', 'n' is stored in f_extra_locals
+        d['m'] = 1
+        d['n'] = 1
+
+        with self.assertRaises(KeyError):
+            d.pop('non_exist')
+
+        del d['m']
+        self.assertEqual(d.pop('n'), 1)
+
+        self.assertNotIn('m', d)
+        self.assertNotIn('n', d)
+
+        self.assertEqual(d.pop('n', 2), 2)
+
     @support.cpython_only
     def test_sizeof(self):
         proxy = sys._getframe().f_locals

Misc/NEWS.d/next/Library/2024-10-16-20-32-40.gh-issue-125590.stHzOP.rst

@@ -0,0 +1 @@
+Allow ``FrameLocalsProxy`` to delete and pop if the key is not a fast variable.

Objects/frameobject.c

@@ -5,6 +5,7 @@
 #include "pycore_code.h"          // CO_FAST_LOCAL, etc.
 #include "pycore_function.h"      // _PyFunction_FromConstructor()
 #include "pycore_moduleobject.h"  // _PyModule_GetDict()
+#include "pycore_modsupport.h"    // _PyArg_CheckPositional()
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include "pycore_opcode_metadata.h" // _PyOpcode_Deopt, _PyOpcode_Caches
 
@@ -158,16 +159,16 @@ framelocalsproxy_setitem(PyObject *self, PyObject *key, PyObject *value)
     PyObject** fast = _PyFrame_GetLocalsArray(frame->f_frame);
     PyCodeObject* co = _PyFrame_GetCode(frame->f_frame);
 
-    if (value == NULL) {
-        PyErr_SetString(PyExc_TypeError, "cannot remove variables from FrameLocalsProxy");
-        return -1;
-    }
-
     int i = framelocalsproxy_getkeyindex(frame, key, false);
     if (i == -2) {
         return -1;
     }
     if (i >= 0) {
+        if (value == NULL) {
+            PyErr_SetString(PyExc_ValueError, "cannot remove local variables from FrameLocalsProxy");
+            return -1;
+        }
+
         _Py_Executors_InvalidateDependency(PyInterpreterState_Get(), co, 1);
 
         _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
@@ -200,6 +201,10 @@ framelocalsproxy_setitem(PyObject *self, PyObject *key, PyObject *value)
     PyObject *extra = frame->f_extra_locals;
 
     if (extra == NULL) {
+        if (value == NULL) {
+            _PyErr_SetKeyError(key);
+            return -1;
+        }
         extra = PyDict_New();
         if (extra == NULL) {
             return -1;
@@ -209,7 +214,11 @@ framelocalsproxy_setitem(PyObject *self, PyObject *key, PyObject *value)
 
     assert(PyDict_Check(extra));
 
-    return PyDict_SetItem(extra, key, value);
+    if (value == NULL) {
+        return PyDict_DelItem(extra, key);
+    } else {
+        return PyDict_SetItem(extra, key, value);
+    }
 }
 
 static int
@@ -674,6 +683,59 @@ framelocalsproxy_setdefault(PyObject* self, PyObject *const *args, Py_ssize_t na
     return result;
 }
 
+static PyObject*
+framelocalsproxy_pop(PyObject* self, PyObject *const *args, Py_ssize_t nargs)
+{
+    if (!_PyArg_CheckPositional("pop", nargs, 1, 2)) {
+        return NULL;
+    }
+
+    PyObject *key = args[0];
+    PyObject *default_value = NULL;
+
+    if (nargs == 2) {
+        default_value = args[1];
+    }
+
+    PyFrameObject *frame = ((PyFrameLocalsProxyObject*)self)->frame;
+
+    int i = framelocalsproxy_getkeyindex(frame, key, false);
+    if (i == -2) {
+        return NULL;
+    }
+
+    if (i >= 0) {
+        PyErr_SetString(PyExc_ValueError, "cannot remove local variables from FrameLocalsProxy");
+        return NULL;
+    }
+
+    PyObject *result = NULL;
+
+    if (frame->f_extra_locals == NULL) {
+        if (default_value != NULL) {
+            return Py_XNewRef(default_value);
+        } else {
+            _PyErr_SetKeyError(key);
+            return NULL;
+        }
+    }
+
+    if (PyDict_Pop(frame->f_extra_locals, key, &result) < 0) {
+        return NULL;
+    }
+
+    if (result == NULL) {
+        if (default_value != NULL) {
+            return Py_XNewRef(default_value);
+        } else {
+            _PyErr_SetKeyError(key);
+            return NULL;
+        }
+    }
+
+    return result;
+}
+
 static PyObject*
 framelocalsproxy_copy(PyObject *self, PyObject *Py_UNUSED(ignored))
 {
@@ -741,6 +803,8 @@ static PyMethodDef framelocalsproxy_methods[] = {
      NULL},
     {"get",          _PyCFunction_CAST(framelocalsproxy_get),            METH_FASTCALL,
      NULL},
+    {"pop",          _PyCFunction_CAST(framelocalsproxy_pop),            METH_FASTCALL,
+     NULL},
     {"setdefault",   _PyCFunction_CAST(framelocalsproxy_setdefault),     METH_FASTCALL,
      NULL},
     {NULL,            NULL}   /* sentinel */