From 0b0d93c320684f9448b6c59f3e13e3ad7fd96d41 Mon Sep 17 00:00:00 2001
From: Petr Viktorin <pviktori@redhat.com>
Date: Tue, 5 May 2020 20:31:49 +0200
Subject: [PATCH] Add _PyType_GetModuleByDef

---
 Include/cpython/object.h           |  2 ++
 Modules/_testmultiphase.c          |  8 ++++++-
 Modules/clinic/_testmultiphase.c.h |  7 ++++--
 Objects/typeobject.c               | 38 ++++++++++++++++++++++++++++++
 4 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/Include/cpython/object.h b/Include/cpython/object.h
index 875a600f79565a..0db53c312f07b7 100644
--- a/Include/cpython/object.h
+++ b/Include/cpython/object.h
@@ -296,6 +296,8 @@ PyAPI_FUNC(PyObject *) _PyObject_LookupSpecial(PyObject *, _Py_Identifier *);
 PyAPI_FUNC(PyTypeObject *) _PyType_CalculateMetaclass(PyTypeObject *, PyObject *);
 PyAPI_FUNC(PyObject *) _PyType_GetDocFromInternalDoc(const char *, const char *);
 PyAPI_FUNC(PyObject *) _PyType_GetTextSignatureFromInternalDoc(const char *, const char *);
+struct PyModuleDef;
+PyAPI_FUNC(PyObject *) _PyType_GetModuleByDef(PyTypeObject *, struct PyModuleDef *);
 
 struct _Py_Identifier;
 PyAPI_FUNC(int) PyObject_Print(PyObject *, FILE *, int);
diff --git a/Modules/_testmultiphase.c b/Modules/_testmultiphase.c
index d69ae628fa7a40..7528422339713f 100644
--- a/Modules/_testmultiphase.c
+++ b/Modules/_testmultiphase.c
@@ -121,24 +121,30 @@ static PyType_Spec Example_Type_spec = {
 };
 
 
+static PyModuleDef def_meth_state_access;
+
 /*[clinic input]
 _testmultiphase.StateAccessType.get_defining_module
 
     cls: defining_class
 
 Return the module of the defining class.
+
+Also tests that result of _PyType_GetModuleByDef matches defining_class's
+module.
 [clinic start generated code]*/
 
 static PyObject *
 _testmultiphase_StateAccessType_get_defining_module_impl(StateAccessTypeObject *self,
                                                          PyTypeObject *cls)
-/*[clinic end generated code: output=ba2a14284a5d0921 input=946149f91cf72c0d]*/
+/*[clinic end generated code: output=ba2a14284a5d0921 input=356f999fc16e0933]*/
 {
     PyObject *retval;
     retval = PyType_GetModule(cls);
     if (retval == NULL) {
         return NULL;
     }
+    assert(_PyType_GetModuleByDef(Py_TYPE(self), &def_meth_state_access) == retval);
     Py_INCREF(retval);
     return retval;
 }
diff --git a/Modules/clinic/_testmultiphase.c.h b/Modules/clinic/_testmultiphase.c.h
index 0d38c230f71865..55f934be8c6c1b 100644
--- a/Modules/clinic/_testmultiphase.c.h
+++ b/Modules/clinic/_testmultiphase.c.h
@@ -6,7 +6,10 @@ PyDoc_STRVAR(_testmultiphase_StateAccessType_get_defining_module__doc__,
 "get_defining_module($self, /)\n"
 "--\n"
 "\n"
-"Return the module of the defining class.");
+"Return the module of the defining class.\n"
+"\n"
+"Also tests that result of _PyType_GetModuleByDef matches defining_class\'s\n"
+"module.");
 
 #define _TESTMULTIPHASE_STATEACCESSTYPE_GET_DEFINING_MODULE_METHODDEF    \
     {"get_defining_module", (PyCFunction)(void(*)(void))_testmultiphase_StateAccessType_get_defining_module, METH_METHOD|METH_FASTCALL|METH_KEYWORDS, _testmultiphase_StateAccessType_get_defining_module__doc__},
@@ -98,4 +101,4 @@ _testmultiphase_StateAccessType_get_count(StateAccessTypeObject *self, PyTypeObj
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=39eea487e94e7f5d input=a9049054013a1b77]*/
+/*[clinic end generated code: output=f01137bb3b373e14 input=a9049054013a1b77]*/
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 36c7662e081a40..a708ceedbf046c 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -3155,6 +3155,44 @@ PyType_GetModuleState(PyTypeObject *type)
     return PyModule_GetState(m);
 }
 
+
+/* Get the module of the first superclass where the module has the
+ * given PyModuleDef.
+ * Implemented by walking the MRO, is relatively slow.
+ *
+ * This is internal API for experimentation within stdlib. Discussion:
+ * https://mail.python.org/archives/list/capi-sig@python.org/thread/T3P2QNLNLBRFHWSKYSTPMVEIL2EEKFJU/
+ */
+PyObject *
+_PyType_GetModuleByDef(PyTypeObject *type, struct PyModuleDef *def)
+{
+    assert(PyType_Check(type));
+    assert(type->tp_mro);
+    int i;
+    for (i = 0; i < PyTuple_GET_SIZE(type->tp_mro); i++) {
+        PyObject *super = PyTuple_GET_ITEM(type->tp_mro, i);
+        if (!PyType_HasFeature((PyTypeObject *)super, Py_TPFLAGS_HEAPTYPE)) {
+            /* Currently, there's no way for static types to inherit
+             * from heap types, but to allow that possibility,
+             * we `continue` rather than `break`.
+             * We'll just potentially loop a few more times before throwing
+             * the error.
+             */
+            continue;
+        }
+        PyHeapTypeObject *ht = (PyHeapTypeObject*)super;
+        if (ht->ht_module && PyModule_GetDef(ht->ht_module) == def) {
+            return ht->ht_module;
+        }
+    }
+    PyErr_Format(
+        PyExc_TypeError,
+        "_PyType_GetModuleByDef: No superclass of '%s' has the given module",
+        type->tp_name);
+    return NULL;
+}
+
+
 /* Internal API to look for a name through the MRO, bypassing the method cache.
    This returns a borrowed reference, and might set an exception.
    'error' is set to: -1: error with exception; 1: error without exception; 0: ok */