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

gaogaotiantian · web-flow · commit afbe5bf9c8c6 · 2024-05-06T20:16:48.000+10:00
Properly implement the `sq_contains` slot for frame locals proxy containment checks.
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
@@ -488,29 +488,33 @@ framelocalsproxy_length(PyObject *self)
     return size;
 }
 
-static PyObject*
+static int
 framelocalsproxy_contains(PyObject *self, PyObject *key)
 {
     PyFrameObject *frame = ((PyFrameLocalsProxyObject*)self)->frame;
 
     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,14 +608,18 @@ 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
     framelocalsproxy_setitem, // mp_ass_subscript
 };
 
 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,
