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

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