Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 28 additions & 2 deletions Lib/test/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.assertRaises(KeyError):
del d['non_exist']

with self.assertRaises(KeyError):
d.pop('non_exist')
gaogaotiantian marked this conversation as resolved.
Show resolved Hide resolved

with self.assertRaises(KeyError):
del d['x']

with self.assertRaises(AttributeError):
d.clear()

with self.assertRaises(AttributeError):
with self.assertRaises(KeyError):
d.pop('x')

with self.assertRaises(KeyError):
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
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Allow ``FrameLocalsProxy`` to delete and pop if the key is not a fast variable.
76 changes: 70 additions & 6 deletions Objects/frameobject.c
Original file line number Diff line number Diff line change
Expand Up @@ -158,16 +158,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_KeyError, "cannot remove local variables from FrameLocalsProxy");
return -1;
}

_Py_Executors_InvalidateDependency(PyInterpreterState_Get(), co, 1);

_PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
Expand Down Expand Up @@ -202,6 +202,10 @@ framelocalsproxy_setitem(PyObject *self, PyObject *key, PyObject *value)
PyObject *extra = frame->f_extra_locals;

if (extra == NULL) {
if (value == NULL) {
PyErr_Format(PyExc_KeyError, "'%R'", key);
return -1;
}
extra = PyDict_New();
if (extra == NULL) {
return -1;
Expand All @@ -211,7 +215,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
Expand Down Expand Up @@ -676,6 +684,60 @@ 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 (nargs < 1 || nargs > 2) {
PyErr_SetString(PyExc_TypeError, "pop expected 1 or 2 arguments");
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_KeyError, "cannot remove local variables from FrameLocalsProxy");
return NULL;
}

PyObject* result = NULL;
gaogaotiantian marked this conversation as resolved.
Show resolved Hide resolved

if (frame->f_extra_locals == NULL) {
if (default_value != NULL) {
return Py_XNewRef(default_value);
} else {
PyErr_Format(PyExc_KeyError, "'%R'", key);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A tiny optimization to consider (though, it's more complex--feel free to reject this):

Suggested change
PyErr_Format(PyExc_KeyError, "'%R'", key);
PyObject *repr = PyObject_Repr(key);
if (repr == NULL) {
return NULL;
}
PyErr_SetObject(PyExc_KeyError, repr);

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just realized I could've used _PyErr_SetKeyError(key) like dictobject.c. I'll switch to that.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, that sounds better. I wasn't aware of that, TIL!

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So I switched to _PyErr_SetKeyError(key) when the key does not exist. However, for local variables cases, I think RuntimeError would fit better. One important reason was KeyError actually expects the "key" to be the explanation string, and giving a reason is a bit weird. Also in this case we have the key, we just can't remove it, so RuntimeError might be better.

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_Format(PyExc_KeyError, "'%R'", key);
return NULL;
}
}

return result;
}

static PyObject*
framelocalsproxy_copy(PyObject *self, PyObject *Py_UNUSED(ignored))
{
Expand Down Expand Up @@ -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 */
Expand Down
Loading