Make sure that tp_dictoffset is correct with Py_TPFLAGS_MANAGED_DICT …

…is set.
python · markshannon · Aug 15, 2022 · Aug 2, 2022 · Aug 2, 2022 · Aug 3, 2022
commit cc5fef4e2f103dbf273d6b93647aa29212883f1d
diff --git a/Lib/test/test_capi.py b/Lib/test/test_capi.py
@@ -538,6 +538,21 @@ def test_heaptype_with_dict(self):
         inst = _testcapi.HeapCTypeWithDict()
         self.assertEqual({}, inst.__dict__)
 
+    def test_heaptype_with_managed_dict(self):
+        inst = _testcapi.HeapCTypeWithManagedDict()
+        inst.foo = 42
+        self.assertEqual(inst.foo, 42)
+        self.assertEqual(inst.__dict__, {"foo": 42})
+
+        inst = _testcapi.HeapCTypeWithManagedDict()
+        self.assertEqual({}, inst.__dict__)
+
+        a = _testcapi.HeapCTypeWithManagedDict()
+        b = _testcapi.HeapCTypeWithManagedDict()
+        a.b = b
+        b.a = a
+        del a, b
+
     def test_heaptype_with_negative_dict(self):
         inst = _testcapi.HeapCTypeWithNegativeDict()
         inst.foo = 42

diff --git a/Modules/_testcapi/heaptype.c b/Modules/_testcapi/heaptype.c
@@ -737,6 +737,45 @@ static PyType_Spec HeapCTypeWithDict_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},
@@ -925,6 +964,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;

@@ -3078,7 +3078,7 @@ type_new_descriptors(const type_new_ctx *ctx, PyTypeObject *type)
     if (ctx->add_dict && ctx->base->tp_itemsize == 0) {
         assert((type->tp_flags & Py_TPFLAGS_MANAGED_DICT) == 0);
         type->tp_flags |= Py_TPFLAGS_MANAGED_DICT;
-        type->tp_dictoffset = -slotoffset - sizeof(PyObject *)*3;
+        type->tp_dictoffset = -type->tp_basicsize - (Py_ssize_t)sizeof(PyObject *)*3;
     }
 
     type->tp_basicsize = slotoffset;
@@ -6078,6 +6078,7 @@ inherit_special(PyTypeObject *type, PyTypeObject *base)
     COPYVAL(tp_itemsize);
     COPYVAL(tp_weaklistoffset);
     COPYVAL(tp_dictoffset);
+
 #undef COPYVAL
 
     /* Setup fast subclass flags */
@@ -6486,6 +6487,22 @@ type_ready_fill_dict(PyTypeObject *type)
     return 0;
 }
 
+static int
+type_ready_dict_offset(PyTypeObject *type)
+{
+    if (type->tp_flags & Py_TPFLAGS_MANAGED_DICT) {
+        if (type->tp_dictoffset > 0) {
+            PyErr_Format(PyExc_TypeError,
+                        "type %s has the Py_TPFLAGS_MANAGED_DICT flag "
+                        "but tp_dictoffset is positive",
+                        type->tp_name);
+            return -1;
+        }
+        Py_ssize_t offset = -type->tp_basicsize - sizeof(PyObject *)*3;
+        type->tp_dictoffset = offset;
+    }
+    return 0;
+}
 
 static int
 type_ready_mro(PyTypeObject *type)
@@ -6694,6 +6711,24 @@ type_ready_post_checks(PyTypeObject *type)
                      type->tp_name);
         return -1;
     }
+    if (type->tp_flags & Py_TPFLAGS_MANAGED_DICT) {
+        if (type->tp_itemsize != 0) {
+            PyErr_Format(PyExc_SystemError,
+                        "type %s has the Py_TPFLAGS_MANAGED_DICT flag "
+                        "but is variable sized",
+                        type->tp_name);
+            return -1;
+        }
+        Py_ssize_t size = type->tp_basicsize;
+        if (type->tp_dictoffset != -size - (Py_ssize_t)sizeof(PyObject *)*3) {
+            PyErr_Format(PyExc_SystemError,
+                        "type %s has the Py_TPFLAGS_MANAGED_DICT flag "
+                        "but tp_dictoffset is set to incompatible value",
+                        type->tp_name);
+            assert(0);
+            return -1;
+        }
+    }
     return 0;
 }
 
@@ -6733,6 +6768,9 @@ type_ready(PyTypeObject *type)
     if (type_ready_inherit(type) < 0) {
         return -1;
     }
+    if (type_ready_dict_offset(type) < 0) {
+        return -1;
+    }
     if (type_ready_set_hash(type) < 0) {
         return -1;
     }