gh-103509: PEP 697 -- Limited C API for Extending Opaque Types (GH-103511)

Co-authored-by: Oleg Iarygin <oleg@arhadthedev.net>
Co-authored-by: Erlend E. Aasland <erlend.aasland@protonmail.com>

5 files changed, 478 insertions, 2 deletions

diff --git a/Modules/Setup.stdlib.in b/Modules/Setup.stdlib.in
index a90c1e9..6b48334 100644
--- a/Modules/Setup.stdlib.in
+++ b/Modules/Setup.stdlib.in
@@ -169,7 +169,7 @@
 @MODULE__XXTESTFUZZ_TRUE@_xxtestfuzz _xxtestfuzz/_xxtestfuzz.c _xxtestfuzz/fuzzer.c
 @MODULE__TESTBUFFER_TRUE@_testbuffer _testbuffer.c
 @MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c
-@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/unicode.c _testcapi/getargs.c _testcapi/pytime.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c _testcapi/pyos.c _testcapi/immortal.c
+@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/unicode.c _testcapi/getargs.c _testcapi/pytime.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c _testcapi/pyos.c _testcapi/immortal.c _testcapi/heaptype_relative.c
 @MODULE__TESTCLINIC_TRUE@_testclinic _testclinic.c
 
 # Some testing modules MUST be built as shared libraries.
diff --git a/Modules/_testcapi/heaptype.c b/Modules/_testcapi/heaptype.c
index 6384fbc..3488e35 100644
--- a/Modules/_testcapi/heaptype.c
+++ b/Modules/_testcapi/heaptype.c
@@ -371,7 +371,6 @@ create_type_from_repeated_slots(PyObject *self, PyObject *variant_obj)
 }
 
 
-
 static PyObject *
 make_immutable_type_with_base(PyObject *self, PyObject *base)
 {
@@ -399,6 +398,17 @@ make_type_with_base(PyObject *self, PyObject *base)
 }
 
 
+static PyObject *
+pyobject_getitemdata(PyObject *self, PyObject *o)
+{
+    void *pointer = PyObject_GetItemData(o);
+    if (pointer == NULL) {
+        return NULL;
+    }
+    return PyLong_FromVoidPtr(pointer);
+}
+
+
 static PyMethodDef TestMethods[] = {
     {"pytype_fromspec_meta",    pytype_fromspec_meta,            METH_O},
     {"test_type_from_ephemeral_spec", test_type_from_ephemeral_spec, METH_NOARGS},
@@ -411,6 +421,7 @@ static PyMethodDef TestMethods[] = {
      METH_NOARGS},
     {"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},
     {NULL},
 };
 
@@ -987,6 +998,113 @@ static PyType_Spec HeapCTypeSetattr_spec = {
     HeapCTypeSetattr_slots
 };
 
+PyDoc_STRVAR(HeapCCollection_doc,
+"Tuple-like heap type that uses PyObject_GetItemData for items.");
+
+static PyObject*
+HeapCCollection_new(PyTypeObject *subtype, PyObject *args, PyObject *kwds)
+{
+    PyObject *self = NULL;
+    PyObject *result = NULL;
+
+    Py_ssize_t size = PyTuple_GET_SIZE(args);
+    self = subtype->tp_alloc(subtype, size);
+    if (!self) {
+        goto finally;
+    }
+    PyObject **data = PyObject_GetItemData(self);
+    if (!data) {
+        goto finally;
+    }
+
+    for (Py_ssize_t i = 0; i < size; i++) {
+        data[i] = Py_NewRef(PyTuple_GET_ITEM(args, i));
+    }
+
+    result = self;
+    self = NULL;
+  finally:
+    Py_XDECREF(self);
+    return result;
+}
+
+static Py_ssize_t
+HeapCCollection_length(PyVarObject *self)
+{
+    return Py_SIZE(self);
+}
+
+static PyObject*
+HeapCCollection_item(PyObject *self, Py_ssize_t i)
+{
+    if (i < 0 || i >= Py_SIZE(self)) {
+        return PyErr_Format(PyExc_IndexError, "index %zd out of range", i);
+    }
+    PyObject **data = PyObject_GetItemData(self);
+    if (!data) {
+        return NULL;
+    }
+    return Py_NewRef(data[i]);
+}
+
+static int
+HeapCCollection_traverse(PyObject *self, visitproc visit, void *arg)
+{
+    PyObject **data = PyObject_GetItemData(self);
+    if (!data) {
+        return -1;
+    }
+    for (Py_ssize_t i = 0; i < Py_SIZE(self); i++) {
+        Py_VISIT(data[i]);
+    }
+    return 0;
+}
+
+static int
+HeapCCollection_clear(PyObject *self)
+{
+    PyObject **data = PyObject_GetItemData(self);
+    if (!data) {
+        return -1;
+    }
+    Py_ssize_t size = Py_SIZE(self);
+    Py_SET_SIZE(self, 0);
+    for (Py_ssize_t i = 0; i < size; i++) {
+        Py_CLEAR(data[i]);
+    }
+    return 0;
+}
+
+static void
+HeapCCollection_dealloc(PyObject *self)
+{
+    PyTypeObject *tp = Py_TYPE(self);
+    HeapCCollection_clear(self);
+    PyObject_GC_UnTrack(self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
+}
+
+static PyType_Slot HeapCCollection_slots[] = {
+    {Py_tp_new, HeapCCollection_new},
+    {Py_sq_length, HeapCCollection_length},
+    {Py_sq_item, HeapCCollection_item},
+    {Py_tp_traverse, HeapCCollection_traverse},
+    {Py_tp_clear, HeapCCollection_clear},
+    {Py_tp_dealloc, HeapCCollection_dealloc},
+    {Py_tp_doc, (void *)HeapCCollection_doc},
+    {0, 0},
+};
+
+static PyType_Spec HeapCCollection_spec = {
+    .name = "_testcapi.HeapCCollection",
+    .basicsize = sizeof(PyVarObject),
+    .itemsize = sizeof(PyObject*),
+    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE |
+              Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_ITEMS_AT_END),
+    .slots = HeapCCollection_slots,
+};
+
 int
 _PyTestCapi_Init_Heaptype(PyObject *m) {
     _testcapimodule = PyModule_GetDef(m);
@@ -1110,5 +1228,16 @@ _PyTestCapi_Init_Heaptype(PyObject *m) {
     }
     PyModule_AddObject(m, "HeapCTypeMetaclassCustomNew", HeapCTypeMetaclassCustomNew);
 
+    PyObject *HeapCCollection = PyType_FromMetaclass(
+        NULL, m, &HeapCCollection_spec, NULL);
+    if (HeapCCollection == NULL) {
+        return -1;
+    }
+    int rc = PyModule_AddType(m, (PyTypeObject *)HeapCCollection);
+    Py_DECREF(HeapCCollection);
+    if (rc < 0) {
+        return -1;
+    }
+
     return 0;
 }
diff --git a/Modules/_testcapi/heaptype_relative.c b/Modules/_testcapi/heaptype_relative.c
new file mode 100644
index 0000000..c247ca3
--- /dev/null
+++ b/Modules/_testcapi/heaptype_relative.c
@@ -0,0 +1,343 @@
+#define Py_LIMITED_API 0x030c0000 // 3.12
+#include "parts.h"
+#include <stddef.h>               // max_align_t
+#include <string.h>               // memset
+
+#ifdef LIMITED_API_AVAILABLE
+
+static PyType_Slot empty_slots[] = {
+    {0, NULL},
+};
+
+static PyObject *
+make_sized_heaptypes(PyObject *module, PyObject *args)
+{
+    PyObject *base = NULL;
+    PyObject *sub = NULL;
+    PyObject *instance = NULL;
+    PyObject *result = NULL;
+
+    int extra_base_size, basicsize;
+
+    int r = PyArg_ParseTuple(args, "ii", &extra_base_size, &basicsize);
+    if (!r) {
+        goto finally;
+    }
+
+    PyType_Spec base_spec = {
+        .name = "_testcapi.Base",
+        .basicsize = sizeof(PyObject) + extra_base_size,
+        .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+        .slots = empty_slots,
+    };
+    PyType_Spec sub_spec = {
+        .name = "_testcapi.Sub",
+        .basicsize = basicsize,
+        .flags = Py_TPFLAGS_DEFAULT,
+        .slots = empty_slots,
+    };
+
+    base = PyType_FromMetaclass(NULL, module, &base_spec, NULL);
+    if (!base) {
+        goto finally;
+    }
+    sub = PyType_FromMetaclass(NULL, module, &sub_spec, base);
+    if (!sub) {
+        goto finally;
+    }
+    instance = PyObject_CallNoArgs(sub);
+    if (!instance) {
+        goto finally;
+    }
+    char *data_ptr = PyObject_GetTypeData(instance, (PyTypeObject *)sub);
+    if (!data_ptr) {
+        goto finally;
+    }
+    Py_ssize_t data_size = PyType_GetTypeDataSize((PyTypeObject *)sub);
+    if (data_size < 0) {
+        goto finally;
+    }
+
+    result = Py_BuildValue("OOOKnn", base, sub, instance,
+                           (unsigned long long)data_ptr,
+                           (Py_ssize_t)(data_ptr - (char*)instance),
+                           data_size);
+  finally:
+    Py_XDECREF(base);
+    Py_XDECREF(sub);
+    Py_XDECREF(instance);
+    return result;
+}
+
+static PyObject *
+var_heaptype_set_data_to_3s(
+    PyObject *self, PyTypeObject *defining_class,
+    PyObject **args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    void *data_ptr = PyObject_GetTypeData(self, defining_class);
+    if (!data_ptr) {
+        return NULL;
+    }
+    Py_ssize_t data_size = PyType_GetTypeDataSize(defining_class);
+    if (data_size < 0) {
+        return NULL;
+    }
+    memset(data_ptr, 3, data_size);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+var_heaptype_get_data(PyObject *self, PyTypeObject *defining_class,
+                      PyObject **args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    void *data_ptr = PyObject_GetTypeData(self, defining_class);
+    if (!data_ptr) {
+        return NULL;
+    }
+    Py_ssize_t data_size = PyType_GetTypeDataSize(defining_class);
+    if (data_size < 0) {
+        return NULL;
+    }
+    return PyBytes_FromStringAndSize(data_ptr, data_size);
+}
+
+static PyMethodDef var_heaptype_methods[] = {
+    {"set_data_to_3s", _PyCFunction_CAST(var_heaptype_set_data_to_3s),
+        METH_METHOD | METH_FASTCALL | METH_KEYWORDS},
+    {"get_data", _PyCFunction_CAST(var_heaptype_get_data),
+        METH_METHOD | METH_FASTCALL | METH_KEYWORDS},
+    {NULL},
+};
+
+static PyObject *
+subclass_var_heaptype(PyObject *module, PyObject *args)
+{
+    PyObject *result = NULL;
+
+    PyObject *base; // borrowed from args
+    int basicsize, itemsize;
+    long pfunc;
+
+    int r = PyArg_ParseTuple(args, "Oiil", &base, &basicsize, &itemsize, &pfunc);
+    if (!r) {
+        goto finally;
+    }
+
+    PyType_Slot slots[] = {
+        {Py_tp_methods, var_heaptype_methods},
+        {0, NULL},
+    };
+
+    PyType_Spec sub_spec = {
+        .name = "_testcapi.Sub",
+        .basicsize = basicsize,
+        .itemsize = itemsize,
+        .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_ITEMS_AT_END,
+        .slots = slots,
+    };
+
+    result = PyType_FromMetaclass(NULL, module, &sub_spec, base);
+  finally:
+    return result;
+}
+
+static PyObject *
+subclass_heaptype(PyObject *module, PyObject *args)
+{
+    PyObject *result = NULL;
+
+    PyObject *base; // borrowed from args
+    int basicsize, itemsize;
+
+    int r = PyArg_ParseTuple(args, "Oii", &base, &basicsize, &itemsize);
+    if (!r) {
+        goto finally;
+    }
+
+    PyType_Slot slots[] = {
+        {Py_tp_methods, var_heaptype_methods},
+        {0, NULL},
+    };
+
+    PyType_Spec sub_spec = {
+        .name = "_testcapi.Sub",
+        .basicsize = basicsize,
+        .itemsize = itemsize,
+        .flags = Py_TPFLAGS_DEFAULT,
+        .slots = slots,
+    };
+
+    result = PyType_FromMetaclass(NULL, module, &sub_spec, base);
+  finally:
+    return result;
+}
+
+static PyMemberDef *
+heaptype_with_member_extract_and_check_memb(PyObject *self)
+{
+    PyMemberDef *def = PyType_GetSlot(Py_TYPE(self), Py_tp_members);
+    if (!def) {
+        if (!PyErr_Occurred()) {
+            PyErr_SetString(PyExc_ValueError, "tp_members is NULL");
+        }
+        return NULL;
+    }
+    if (!def[0].name) {
+        PyErr_SetString(PyExc_ValueError, "tp_members[0] is NULL");
+        return NULL;
+    }
+    if (def[1].name) {
+        PyErr_SetString(PyExc_ValueError, "tp_members[1] is not NULL");
+        return NULL;
+    }
+    if (strcmp(def[0].name, "memb")) {
+        PyErr_SetString(PyExc_ValueError, "tp_members[0] is not for `memb`");
+        return NULL;
+    }
+    if (def[0].flags) {
+        PyErr_SetString(PyExc_ValueError, "tp_members[0] has flags set");
+        return NULL;
+    }
+    return def;
+}
+
+static PyObject *
+heaptype_with_member_get_memb(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    PyMemberDef *def = heaptype_with_member_extract_and_check_memb(self);
+    return PyMember_GetOne((const char *)self, def);
+}
+
+static PyObject *
+heaptype_with_member_set_memb(PyObject *self, PyObject *value)
+{
+    PyMemberDef *def = heaptype_with_member_extract_and_check_memb(self);
+    int r = PyMember_SetOne((char *)self, def, value);
+    if (r < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+get_memb_offset(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    PyMemberDef *def = heaptype_with_member_extract_and_check_memb(self);
+    return PyLong_FromSsize_t(def->offset);
+}
+
+static PyObject *
+heaptype_with_member_get_memb_relative(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    PyMemberDef def = {"memb", Py_T_BYTE, sizeof(PyObject), Py_RELATIVE_OFFSET};
+    return PyMember_GetOne((const char *)self, &def);
+}
+
+static PyObject *
+heaptype_with_member_set_memb_relative(PyObject *self, PyObject *value)
+{
+    PyMemberDef def = {"memb", Py_T_BYTE, sizeof(PyObject), Py_RELATIVE_OFFSET};
+    int r = PyMember_SetOne((char *)self, &def, value);
+    if (r < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+static PyMethodDef heaptype_with_member_methods[] = {
+    {"get_memb", heaptype_with_member_get_memb, METH_NOARGS},
+    {"set_memb", heaptype_with_member_set_memb, METH_O},
+    {"get_memb_offset", get_memb_offset, METH_NOARGS},
+    {"get_memb_relative", heaptype_with_member_get_memb_relative, METH_NOARGS},
+    {"set_memb_relative", heaptype_with_member_set_memb_relative, METH_O},
+    {NULL},
+};
+
+static PyObject *
+make_heaptype_with_member(PyObject *module, PyObject *args)
+{
+    PyObject *base = NULL;
+    PyObject *result = NULL;
+
+    int extra_base_size, basicsize, offset, add_flag;
+
+    int r = PyArg_ParseTuple(args, "iiip", &extra_base_size, &basicsize, &offset, &add_flag);
+    if (!r) {
+        goto finally;
+    }
+
+    PyType_Spec base_spec = {
+        .name = "_testcapi.Base",
+        .basicsize = sizeof(PyObject) + extra_base_size,
+        .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
+        .slots = empty_slots,
+    };
+    base = PyType_FromMetaclass(NULL, module, &base_spec, NULL);
+    if (!base) {
+        goto finally;
+    }
+
+    PyMemberDef members[] = {
+        {"memb", Py_T_BYTE, offset, add_flag ? Py_RELATIVE_OFFSET : 0},
+        {0},
+    };
+    PyType_Slot slots[] = {
+        {Py_tp_members, members},
+        {Py_tp_methods, heaptype_with_member_methods},
+        {0, NULL},
+    };
+
+    PyType_Spec sub_spec = {
+        .name = "_testcapi.Sub",
+        .basicsize = basicsize,
+        .flags = Py_TPFLAGS_DEFAULT,
+        .slots = slots,
+    };
+
+    result = PyType_FromMetaclass(NULL, module, &sub_spec, base);
+  finally:
+    Py_XDECREF(base);
+    return result;
+}
+
+
+static PyObject *
+test_alignof_max_align_t(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    // We define ALIGNOF_MAX_ALIGN_T even if the compiler doesn't support
+    // max_align_t. Double-check that it's correct.
+    assert(ALIGNOF_MAX_ALIGN_T > 0);
+    assert(ALIGNOF_MAX_ALIGN_T >= _Alignof(long long));
+    assert(ALIGNOF_MAX_ALIGN_T >= _Alignof(long double));
+    assert(ALIGNOF_MAX_ALIGN_T >= _Alignof(void*));
+    assert(ALIGNOF_MAX_ALIGN_T >= _Alignof(void (*)(void)));
+
+    // Ensure it's a power of two
+    assert((ALIGNOF_MAX_ALIGN_T & (ALIGNOF_MAX_ALIGN_T - 1)) == 0);
+
+    Py_RETURN_NONE;
+}
+
+static PyMethodDef TestMethods[] = {
+    {"make_sized_heaptypes", make_sized_heaptypes, METH_VARARGS},
+    {"subclass_var_heaptype", subclass_var_heaptype, METH_VARARGS},
+    {"subclass_heaptype", subclass_heaptype, METH_VARARGS},
+    {"make_heaptype_with_member", make_heaptype_with_member, METH_VARARGS},
+    {"test_alignof_max_align_t", test_alignof_max_align_t, METH_NOARGS},
+    {NULL},
+};
+
+int
+_PyTestCapi_Init_HeaptypeRelative(PyObject *m) {
+    if (PyModule_AddFunctions(m, TestMethods) < 0) {
+        return -1;
+    }
+
+    if (PyModule_AddIntMacro(m, ALIGNOF_MAX_ALIGN_T) < 0) {
+        return -1;
+    }
+
+    return 0;
+}
+
+#endif // LIMITED_API_AVAILABLE
diff --git a/Modules/_testcapi/parts.h b/Modules/_testcapi/parts.h
index 4d2d683..d75412d 100644
--- a/Modules/_testcapi/parts.h
+++ b/Modules/_testcapi/parts.h
@@ -43,6 +43,7 @@ int _PyTestCapi_Init_Immortal(PyObject *module);
 
 #ifdef LIMITED_API_AVAILABLE
 int _PyTestCapi_Init_VectorcallLimited(PyObject *module);
+int _PyTestCapi_Init_HeaptypeRelative(PyObject *module);
 #endif // LIMITED_API_AVAILABLE
 
 #endif
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 47e0ed9..1ecc442 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -4324,6 +4324,9 @@ PyInit__testcapi(void)
     if (_PyTestCapi_Init_VectorcallLimited(m) < 0) {
         return NULL;
     }
+    if (_PyTestCapi_Init_HeaptypeRelative(m) < 0) {
+        return NULL;
+    }
 #endif
 
     PyState_AddModule(m, &_testcapimodule);