gh-74929: Make containment checks more efficient in `FrameLocalsProxy` (#118624)

Properly implement the `sq_contains` slot for frame locals proxy containment checks.
This commit is contained in:
Tian Gao 2024-05-06 03:16:48 -07:00 committed by GitHub
parent 757fd3e010
commit afbe5bf9c8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 19 additions and 10 deletions

View File

@ -488,7 +488,7 @@ framelocalsproxy_length(PyObject *self)
return size;
}
static PyObject*
static int
framelocalsproxy_contains(PyObject *self, PyObject *key)
{
PyFrameObject *frame = ((PyFrameLocalsProxyObject*)self)->frame;
@ -496,21 +496,25 @@ framelocalsproxy_contains(PyObject *self, PyObject *key)
if (PyUnicode_CheckExact(key)) {
int i = framelocalsproxy_getkeyindex(frame, key, true);
if (i >= 0) {
Py_RETURN_TRUE;
return 1;
}
}
PyObject *extra = ((PyFrameObject*)frame)->f_extra_locals;
if (extra != NULL) {
int result = PyDict_Contains(extra, key);
if (result < 0) {
return NULL;
} else if (result > 0) {
Py_RETURN_TRUE;
}
return PyDict_Contains(extra, key);
}
Py_RETURN_FALSE;
return 0;
}
static PyObject* framelocalsproxy___contains__(PyObject *self, PyObject *key)
{
int result = framelocalsproxy_contains(self, key);
if (result < 0) {
return NULL;
}
return PyBool_FromLong(result);
}
static PyObject*
@ -604,6 +608,10 @@ static PyNumberMethods framelocalsproxy_as_number = {
.nb_inplace_or = framelocalsproxy_inplace_or,
};
static PySequenceMethods framelocalsproxy_as_sequence = {
.sq_contains = framelocalsproxy_contains,
};
static PyMappingMethods framelocalsproxy_as_mapping = {
framelocalsproxy_length, // mp_length
framelocalsproxy_getitem, // mp_subscript
@ -611,7 +619,7 @@ static PyMappingMethods framelocalsproxy_as_mapping = {
};
static PyMethodDef framelocalsproxy_methods[] = {
{"__contains__", framelocalsproxy_contains, METH_O | METH_COEXIST,
{"__contains__", framelocalsproxy___contains__, METH_O | METH_COEXIST,
NULL},
{"__getitem__", framelocalsproxy_getitem, METH_O | METH_COEXIST,
NULL},
@ -639,6 +647,7 @@ PyTypeObject PyFrameLocalsProxy_Type = {
.tp_dealloc = (destructor)framelocalsproxy_dealloc,
.tp_repr = &framelocalsproxy_repr,
.tp_as_number = &framelocalsproxy_as_number,
.tp_as_sequence = &framelocalsproxy_as_sequence,
.tp_as_mapping = &framelocalsproxy_as_mapping,
.tp_getattro = PyObject_GenericGetAttr,
.tp_setattro = PyObject_GenericSetAttr,