GH-95707: Fix uses of `Py_TPFLAGS_MANAGED_DICT` (GH-95854)

* Make sure that tp_dictoffset is correct with Py_TPFLAGS_MANAGED_DICT is set.

* Avoid traversing managed dict twice when subclassing class with Py_TPFLAGS_MANAGED_DICT set.

1 files changed, 45 insertions, 0 deletions

diff --git a/Modules/_testcapi/heaptype.c b/Modules/_testcapi/heaptype.c
index 514541c..de990ba 100644
--- a/Modules/_testcapi/heaptype.c
+++ b/Modules/_testcapi/heaptype.c
@@ -761,6 +761,45 @@ static PyType_Spec HeapCTypeWithDict2_spec = {
     HeapCTypeWithDict_slots
 };
 
+static int
+heapmanaged_traverse(HeapCTypeObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(Py_TYPE(self));
+    return _PyObject_VisitManagedDict((PyObject *)self, visit, arg);
+}
+
+static void
+heapmanaged_clear(HeapCTypeObject *self)
+{
+    _PyObject_ClearManagedDict((PyObject *)self);
+}
+
+static void
+heapmanaged_dealloc(HeapCTypeObject *self)
+{
+    PyTypeObject *tp = Py_TYPE(self);
+    _PyObject_ClearManagedDict((PyObject *)self);
+    PyObject_GC_UnTrack(self);
+    PyObject_GC_Del(self);
+    Py_DECREF(tp);
+}
+
+static PyType_Slot HeapCTypeWithManagedDict_slots[] = {
+    {Py_tp_traverse, heapmanaged_traverse},
+    {Py_tp_getset, heapctypewithdict_getsetlist},
+    {Py_tp_clear, heapmanaged_clear},
+    {Py_tp_dealloc, heapmanaged_dealloc},
+    {0, 0},
+};
+
+static PyType_Spec  HeapCTypeWithManagedDict_spec = {
+    "_testcapi.HeapCTypeWithManagedDict",
+    sizeof(PyObject),
+    0,
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_MANAGED_DICT,
+    HeapCTypeWithManagedDict_slots
+};
+
 static struct PyMemberDef heapctypewithnegativedict_members[] = {
     {"dictobj", T_OBJECT, offsetof(HeapCTypeWithDictObject, dict)},
     {"__dictoffset__", T_PYSSIZET, -(Py_ssize_t)sizeof(void*), READONLY},
@@ -963,6 +1002,12 @@ _PyTestCapi_Init_Heaptype(PyObject *m) {
     }
     PyModule_AddObject(m, "HeapCTypeWithNegativeDict", HeapCTypeWithNegativeDict);
 
+    PyObject *HeapCTypeWithManagedDict = PyType_FromSpec(&HeapCTypeWithManagedDict_spec);
+    if (HeapCTypeWithManagedDict == NULL) {
+        return -1;
+    }
+    PyModule_AddObject(m, "HeapCTypeWithManagedDict", HeapCTypeWithManagedDict);
+
     PyObject *HeapCTypeWithWeakref = PyType_FromSpec(&HeapCTypeWithWeakref_spec);
     if (HeapCTypeWithWeakref == NULL) {
         return -1;