expose _PyObject_LookupSpecialMethod to types.lookup_special_method

Author: tanloong

Doc/library/types.rst

@@ -521,6 +521,14 @@ Additional Utility Classes and Functions
 
    .. versionadded:: 3.4
 
+.. function:: lookup_special_method(obj, attr)
+
+   Do a method lookup in the type without looking in the instance dictionary
+   but still binding it to the instance. Returns None if the method is not
+   found.
+
+   .. versionadded:: 3.15
+
 
 Coroutine Utility Functions
 ---------------------------

Lib/test/test_types.py

@@ -41,7 +41,8 @@ def clear_typing_caches():
 class TypesTests(unittest.TestCase):
 
     def test_names(self):
-        c_only_names = {'CapsuleType', 'LazyImportType'}
+        c_only_names = {'CapsuleType', 'LazyImportType',
+                        'lookup_special_method'}
         ignored = {'new_class', 'resolve_bases', 'prepare_class',
                    'get_original_bases', 'DynamicClassAttribute', 'coroutine'}
 
@@ -59,7 +60,7 @@ def test_names(self):
             'MemberDescriptorType', 'MethodDescriptorType', 'MethodType',
             'MethodWrapperType', 'ModuleType', 'NoneType',
             'NotImplementedType', 'SimpleNamespace', 'TracebackType',
-            'UnionType', 'WrapperDescriptorType',
+            'UnionType', 'WrapperDescriptorType', 'lookup_special_method',
         }
         self.assertEqual(all_names, set(c_types.__all__))
         self.assertEqual(all_names - c_only_names, set(py_types.__all__))
@@ -726,6 +727,46 @@ def test_frame_locals_proxy_type(self):
         self.assertIsNotNone(frame)
         self.assertIsInstance(frame.f_locals, types.FrameLocalsProxyType)
 
+    def test_lookup_special_method(self):
+        class CM1:
+            def __enter__(self):
+                return "__enter__ from class __dict__"
+
+        class CM2:
+            def __init__(self):
+                def __enter__(self):
+                    return "__enter__ from instance __dict__"
+                self.__enter__ = __enter__
+
+        class CM3:
+            __slots__ = ("__enter__",)
+            def __init__(self):
+                def __enter__(self):
+                    return "__enter__ from __slots__"
+                self.__enter__ = __enter__
+        cm1 = CM1()
+        meth = types.lookup_special_method(cm1, "__enter__")
+        self.assertIsNotNone(meth)
+        self.assertEqual(meth(cm1), "__enter__ from class __dict__")
+
+        meth = types.lookup_special_method(cm1, "__missing__")
+        self.assertIsNone(meth)
+
+        with self.assertRaises(TypeError):
+            types.lookup_special_method(cm1, 123)
+
+        cm2 = CM2()
+        meth = types.lookup_special_method(cm2, "__enter__")
+        self.assertIsNone(meth)
+
+        cm3 = CM3()
+        meth = types.lookup_special_method(cm3, "__enter__")
+        self.assertIsNotNone(meth)
+        self.assertEqual(meth(cm3), "__enter__ from __slots__")
+
+        meth = types.lookup_special_method([], "__len__")
+        self.assertIsNotNone(meth)
+        self.assertEqual(meth([]), 0)
 
 class UnionTests(unittest.TestCase):
 

Modules/_typesmodule.c

@@ -6,6 +6,34 @@
 #include "pycore_namespace.h"     // _PyNamespace_Type
 #include "pycore_object.h"        // _PyNone_Type, _PyNotImplemented_Type
 #include "pycore_unionobject.h"   // _PyUnion_Type
+#include "pycore_typeobject.h"    // _PyObject_LookupSpecialMethod
+#include "pycore_stackref.h"      // _PyStackRef
+
+static PyObject *
+_types_lookup_special_method_impl(PyObject *self, PyObject *args)
+{
+  PyObject *obj, *attr;
+  if (!PyArg_ParseTuple(args, "OO", &obj, &attr)) {
+    return NULL;
+  }
+  if (!PyUnicode_Check(attr)) {
+    PyErr_Format(PyExc_TypeError, "attribute name must be string, not '%.200s'",
+                 Py_TYPE(attr)->tp_name);
+    return NULL;
+  }
+  _PyStackRef method_and_self[2];
+  method_and_self[0] = PyStackRef_NULL;
+  method_and_self[1] = PyStackRef_FromPyObjectBorrow(obj);
+  int result = _PyObject_LookupSpecialMethod(attr, method_and_self);
+  if (result == -1) {
+    return NULL;
+  }
+  if (result == 0) {
+    Py_RETURN_NONE;
+  }
+  PyObject *method = PyStackRef_AsPyObjectSteal(method_and_self[0]);
+  return Py_BuildValue("O", method);
+}
 
 static int
 _types_exec(PyObject *m)
@@ -60,12 +88,20 @@ static struct PyModuleDef_Slot _typesmodule_slots[] = {
     {0, NULL}
 };
 
+static PyMethodDef _typesmodule_methods[] = {
+    {"lookup_special_method", _types_lookup_special_method_impl, METH_VARARGS,
+     "Do a method lookup in the type without looking in the instance "
+     "dictionary but still binding it to the instance. Returns None if the "
+     "method is not found."},
+    {NULL, NULL, 0, NULL}};
+
 static struct PyModuleDef typesmodule = {
     .m_base = PyModuleDef_HEAD_INIT,
     .m_name = "_types",
     .m_doc = "Define names for built-in types.",
     .m_size = 0,
     .m_slots = _typesmodule_slots,
+    .m_methods = _typesmodule_methods,
 };
 
 PyMODINIT_FUNC