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

This commit is contained in:
Tian Gao 2024-10-21 08:43:08 -07:00 committed by GitHub
parent 3d1df3d84e
commit 5b7a872b26
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 99 additions and 8 deletions

View File

@ -397,15 +397,41 @@ class TestFrameLocals(unittest.TestCase):
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

View File

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

View File

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