gh-74929: Make containment checks more efficient in `FrameLocalsProxy` (#118624)
Properly implement the `sq_contains` slot for frame locals proxy containment checks.
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index fdac869..aac809b 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -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,
