From 79bd50069f1d173a279876f12c76fc75224a00a5 Mon Sep 17 00:00:00 2001
From: Victor Stinner <vstinner@python.org>
Date: Tue, 3 Oct 2023 17:27:51 +0200
Subject: [PATCH] Add PyObject_VisitManagedDict()

Add PyObject_VisitManagedDict() and PyObject_ClearManagedDict()
functions.
---
 docs/api.rst                        |  8 +++
 docs/changelog.rst                  |  2 +
 pythoncapi_compat.h                 | 25 +++++++++
 tests/test_pythoncapi_compat_cext.c | 82 +++++++++++++++++++++++++++++
 4 files changed, 117 insertions(+)

diff --git a/docs/api.rst b/docs/api.rst
index c0adfc4..c2564f5 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -99,6 +99,14 @@ Python 3.13
 
    See `PyLong_AsInt() documentation <https://docs.python.org/dev/c-api/long.html#c.PyLong_AsInt>`__.
 
+.. c:function:: int PyObject_VisitManagedDict(PyObject *obj, visitproc visit, void *arg)
+
+   See `PyObject_VisitManagedDict() documentation <https://docs.python.org/dev/c-api/object.html#c.PyObject_VisitManagedDict>`__.
+
+.. c:function:: void PyObject_ClearManagedDict(PyObject *obj)
+
+   See `PyObject_ClearManagedDict() documentation <https://docs.python.org/dev/c-api/object.html#c.PyObject_ClearManagedDict>`__.
+
 
 Python 3.12
 -----------
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 7a21348..b5203b8 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -1,6 +1,8 @@
 Changelog
 =========
 
+* 2023-10-03: Add ``PyObject_VisitManagedDict()`` and
+  ``PyObject_ClearManagedDict()`` functions.
 * 2023-09-29: Add functions:
 
   * ``PyMapping_HasKeyWithError()``
diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h
index e6a388e..fa32887 100644
--- a/pythoncapi_compat.h
+++ b/pythoncapi_compat.h
@@ -905,6 +905,31 @@ static inline int PyLong_AsInt(PyObject *obj)
 #endif
 
 
+// gh-107073 added PyObject_VisitManagedDict() to Python 3.13.0a1
+#if PY_VERSION_HEX < 0x030D00A1
+static inline int
+PyObject_VisitManagedDict(PyObject *obj, visitproc visit, void *arg)
+{
+    PyObject **dict = _PyObject_GetDictPtr(obj);
+    if (*dict == NULL) {
+        return -1;
+    }
+    Py_VISIT(*dict);
+    return 0;
+}
+
+static inline void
+PyObject_ClearManagedDict(PyObject *obj)
+{
+    PyObject **dict = _PyObject_GetDictPtr(obj);
+    if (*dict == NULL) {
+        return;
+    }
+    Py_CLEAR(*dict);
+}
+#endif
+
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/tests/test_pythoncapi_compat_cext.c b/tests/test_pythoncapi_compat_cext.c
index 7e08402..b50333d 100644
--- a/tests/test_pythoncapi_compat_cext.c
+++ b/tests/test_pythoncapi_compat_cext.c
@@ -1277,6 +1277,85 @@ test_long_api(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
 }
 
 
+// --- HeapCTypeWithManagedDict --------------------------------------------
+
+// Py_TPFLAGS_MANAGED_DICT was added to Python 3.11.0a3
+#if PY_VERSION_HEX >= 0x030B00A3
+#  define TEST_MANAGED_DICT
+
+typedef struct {
+    PyObject_HEAD
+} HeapCTypeObject;
+
+static int
+heapmanaged_traverse(PyObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(Py_TYPE(self));
+    return PyObject_VisitManagedDict(self, visit, arg);
+}
+
+static int
+heapmanaged_clear(PyObject *self)
+{
+    PyObject_ClearManagedDict(self);
+    return 0;
+}
+
+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, _Py_CAST(void*, heapmanaged_traverse)},
+    {Py_tp_clear, _Py_CAST(void*, heapmanaged_clear)},
+    {Py_tp_dealloc, _Py_CAST(void*, heapmanaged_dealloc)},
+    {0, 0},
+};
+
+static PyType_Spec HeapCTypeWithManagedDict_spec = {
+    "test_pythoncapi_compat.HeapCTypeWithManagedDict",
+    sizeof(PyObject),
+    0,
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_MANAGED_DICT,
+    HeapCTypeWithManagedDict_slots
+};
+
+static PyObject *
+test_managed_dict(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
+{
+    PyObject *type = PyType_FromSpec(&HeapCTypeWithManagedDict_spec);
+    if (type == NULL) {
+        return NULL;
+    }
+
+    PyObject *obj = PyObject_CallNoArgs(type);
+    if (obj == NULL) {
+        Py_DECREF(type);
+        return NULL;
+    }
+
+    // call heapmanaged_traverse()
+    PyGC_Collect();
+
+    // call heapmanaged_clear()
+    Py_DECREF(obj);
+    PyGC_Collect();
+
+    Py_DECREF(type);
+    // Just in case!
+    PyGC_Collect();
+
+    Py_RETURN_NONE;
+}
+#endif  // PY_VERSION_HEX >= 0x030B00A3
+
+
 static struct PyMethodDef methods[] = {
     {"test_object", test_object, METH_NOARGS, _Py_NULL},
     {"test_py_is", test_py_is, METH_NOARGS, _Py_NULL},
@@ -1303,6 +1382,9 @@ static struct PyMethodDef methods[] = {
     {"test_getitem", test_getitem, METH_NOARGS, _Py_NULL},
     {"test_dict_api", test_dict_api, METH_NOARGS, _Py_NULL},
     {"test_long_api", test_long_api, METH_NOARGS, _Py_NULL},
+#ifdef TEST_MANAGED_DICT
+    {"test_managed_dict", test_managed_dict, METH_NOARGS, _Py_NULL},
+#endif
     {_Py_NULL, _Py_NULL, 0, _Py_NULL}
 };