gh-124153: Implement PyType_GetBaseByToken() and Py_tp_token slot (GH-124163)

Modules/_ctypes/_ctypes.c

@@ -500,7 +500,7 @@ CType_Type_dealloc(PyObject *self)
 {
     StgInfo *info = _PyStgInfo_FromType_NoState(self);
     if (!info) {
-        PyErr_WriteUnraisable(self);
+        PyErr_WriteUnraisable(NULL);  // NULL avoids segfault here
     }
     if (info) {
         PyMem_Free(info->ffi_type_pointer.elements);
@@ -560,6 +560,7 @@ static PyMethodDef ctype_methods[] = {
 };
 
 static PyType_Slot ctype_type_slots[] = {
+    {Py_tp_token, Py_TP_USE_SPEC},
     {Py_tp_traverse, CType_Type_traverse},
     {Py_tp_clear, CType_Type_clear},
     {Py_tp_dealloc, CType_Type_dealloc},
@@ -569,7 +570,7 @@ static PyType_Slot ctype_type_slots[] = {
     {0, NULL},
 };
 
-static PyType_Spec pyctype_type_spec = {
+PyType_Spec pyctype_type_spec = {
     .name = "_ctypes.CType_Type",
     .basicsize = -(Py_ssize_t)sizeof(StgInfo),
     .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE |

Modules/_ctypes/ctypes.h

@@ -108,6 +108,7 @@ get_module_state_by_def(PyTypeObject *cls)
 }
 
 
+extern PyType_Spec pyctype_type_spec;
 extern PyType_Spec carg_spec;
 extern PyType_Spec cfield_spec;
 extern PyType_Spec cthunk_spec;
@@ -490,16 +491,23 @@ PyStgInfo_FromAny(ctypes_state *state, PyObject *obj, StgInfo **result)
 
 /* A variant of PyStgInfo_FromType that doesn't need the state,
  * so it can be called from finalization functions when the module
- * state is torn down. Does no checks; cannot fail.
- * This inlines the current implementation PyObject_GetTypeData,
- * so it might break in the future.
+ * state is torn down.
  */
 static inline StgInfo *
 _PyStgInfo_FromType_NoState(PyObject *type)
 {
-    size_t type_basicsize =_Py_SIZE_ROUND_UP(PyType_Type.tp_basicsize,
-                                             ALIGNOF_MAX_ALIGN_T);
-    return (StgInfo *)((char *)type + type_basicsize);
+    PyTypeObject *PyCType_Type;
+    if (PyType_GetBaseByToken(Py_TYPE(type), &pyctype_type_spec, &PyCType_Type) < 0) {
+        return NULL;
+    }
+    if (PyCType_Type == NULL) {
+        PyErr_Format(PyExc_TypeError, "expected a ctypes type, got '%N'", type);
+        return NULL;
+    }
+
+    StgInfo *info = PyObject_GetTypeData(type, PyCType_Type);
+    Py_DECREF(PyCType_Type);
+    return info;
 }
 
 // Initialize StgInfo on a newly created type

Modules/_testcapi/heaptype.c

@@ -410,6 +410,118 @@ pyobject_getitemdata(PyObject *self, PyObject *o)
 }
 
 
+static PyObject *
+create_type_with_token(PyObject *module, PyObject *args)
+{
+    const char *name;
+    PyObject *py_token;
+    if (!PyArg_ParseTuple(args, "sO", &name, &py_token)) {
+        return NULL;
+    }
+    void *token = PyLong_AsVoidPtr(py_token);
+    if (token == Py_TP_USE_SPEC) {
+        // Py_TP_USE_SPEC requires the spec that at least outlives the class
+        static PyType_Slot slots[] = {
+            {Py_tp_token, Py_TP_USE_SPEC},
+            {0},
+        };
+        static PyType_Spec spec = {
+            .name = "_testcapi.DefaultTokenTest",
+            .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+            .slots = slots,
+        };
+        PyObject *type = PyType_FromMetaclass(NULL, NULL, &spec, NULL);
+        if (!type) {
+            return NULL;
+        }
+        token = PyType_GetSlot((PyTypeObject *)type, Py_tp_token);
+        assert(!PyErr_Occurred());
+        Py_DECREF(type);
+        if (token != &spec) {
+            PyErr_SetString(PyExc_AssertionError,
+                            "failed to convert token from Py_TP_USE_SPEC");
+            return NULL;
+        }
+    }
+    // Test non-NULL token that must also outlive the class
+    PyType_Slot slots[] = {
+        {Py_tp_token, token},
+        {0},
+    };
+    PyType_Spec spec = {
+        .name = name,
+        .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+        .slots = slots,
+    };
+    return PyType_FromMetaclass(NULL, module, &spec, NULL);
+}
+
+static PyObject *
+get_tp_token(PyObject *self, PyObject *type)
+{
+    void *token = PyType_GetSlot((PyTypeObject *)type, Py_tp_token);
+    if (PyErr_Occurred()) {
+        return NULL;
+    }
+    return PyLong_FromVoidPtr(token);
+}
+
+static PyObject *
+pytype_getbasebytoken(PyObject *self, PyObject *args)
+{
+    PyTypeObject *type;
+    PyObject *py_token, *use_mro, *need_result;
+    if (!PyArg_ParseTuple(args, "OOOO",
+                          &type, &py_token, &use_mro, &need_result)) {
+        return NULL;
+    }
+
+    PyObject *mro_save = NULL;
+    if (use_mro != Py_True) {
+        // Test internal detail: PyType_GetBaseByToken works even with
+        // types that are only partially initialized (or torn down):
+        // if tp_mro=NULL we fall back to tp_bases.
+        assert(PyType_Check(type));
+        mro_save = type->tp_mro;
+        type->tp_mro = NULL;
+    }
+
+    void *token = PyLong_AsVoidPtr(py_token);
+    PyObject *result;
+    int ret;
+    if (need_result == Py_True) {
+        ret = PyType_GetBaseByToken(type, token, (PyTypeObject **)&result);
+    }
+    else {
+        result = NULL;
+        ret = PyType_GetBaseByToken(type, token, NULL);
+    }
+
+    if (use_mro != Py_True) {
+        type->tp_mro = mro_save;
+    }
+    if (ret < 0) {
+        assert(result == NULL);
+        return NULL;
+    }
+    PyObject *py_ret = PyLong_FromLong(ret);
+    if (py_ret == NULL) {
+        goto error;
+    }
+    PyObject *tuple = PyTuple_New(2);
+    if (tuple == NULL) {
+        goto error;
+    }
+    PyTuple_SET_ITEM(tuple, 0, py_ret);
+    PyTuple_SET_ITEM(tuple, 1, result ? result : Py_None);
+    return tuple;
+error:
+    Py_XDECREF(py_ret);
+    Py_XDECREF(result);
+    return NULL;
+}
+
+
 static PyMethodDef TestMethods[] = {
     {"pytype_fromspec_meta",    pytype_fromspec_meta,            METH_O},
     {"test_type_from_ephemeral_spec", test_type_from_ephemeral_spec, METH_NOARGS},
@@ -423,6 +535,9 @@ static PyMethodDef TestMethods[] = {
     {"make_immutable_type_with_base", make_immutable_type_with_base, METH_O},
     {"make_type_with_base", make_type_with_base, METH_O},
     {"pyobject_getitemdata", pyobject_getitemdata, METH_O},
+    {"create_type_with_token", create_type_with_token, METH_VARARGS},
+    {"get_tp_token", get_tp_token, METH_O},
+    {"pytype_getbasebytoken", pytype_getbasebytoken, METH_VARARGS},
     {NULL},
 };
 
@@ -1287,6 +1402,8 @@ _PyTestCapi_Init_Heaptype(PyObject *m) {
         &PyType_Type, m, &HeapCTypeMetaclassNullNew_spec, (PyObject *) &PyType_Type);
     ADD("HeapCTypeMetaclassNullNew", HeapCTypeMetaclassNullNew);
 
+    ADD("Py_TP_USE_SPEC", PyLong_FromVoidPtr(Py_TP_USE_SPEC));
+
     PyObject *HeapCCollection = PyType_FromMetaclass(
         NULL, m, &HeapCCollection_spec, NULL);
     if (HeapCCollection == NULL) {