[3.13] gh-125590: Allow FrameLocalsProxy to delete and pop keys from extra locals (GH-125616)

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 */