From 5b7a872b26a9ba6c93d7c2109559a82d1c1612de Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Mon, 21 Oct 2024 08:43:08 -0700 Subject: [PATCH] gh-125590: Allow FrameLocalsProxy to delete and pop keys from extra locals (#125616) --- Lib/test/test_frame.py | 30 +++++++- ...-10-16-20-32-40.gh-issue-125590.stHzOP.rst | 1 + Objects/frameobject.c | 76 +++++++++++++++++-- 3 files changed, 99 insertions(+), 8 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2024-10-16-20-32-40.gh-issue-125590.stHzOP.rst diff --git a/Lib/test/test_frame.py b/Lib/test/test_frame.py index 32de8ed9a13f80..11f191700ccef0 100644 --- a/Lib/test/test_frame.py +++ b/Lib/test/test_frame.py @@ -397,15 +397,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 diff --git a/Misc/NEWS.d/next/Library/2024-10-16-20-32-40.gh-issue-125590.stHzOP.rst b/Misc/NEWS.d/next/Library/2024-10-16-20-32-40.gh-issue-125590.stHzOP.rst new file mode 100644 index 00000000000000..dc6765ada641a9 --- /dev/null +++ b/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. diff --git a/Objects/frameobject.c b/Objects/frameobject.c index f3a66ffc9aac8f..5ef48919a081be 100644 --- a/Objects/frameobject.c +++ b/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) _PyStackRef *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); @@ -202,6 +203,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; @@ -211,7 +216,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 @@ -676,6 +685,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)) { @@ -743,6 +805,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 */