bpo-40077: Convert zoneinfo type to heap type

python · erlend-aasland · Feb 6, 2021 · Feb 8, 2021 · Feb 8, 2021 · Feb 8, 2021
commit 073a712f38d9ce4debf597e63f1c0139481c146f
diff --git a/Lib/test/test_zoneinfo/test_zoneinfo.py b/Lib/test/test_zoneinfo/test_zoneinfo.py
@@ -1775,12 +1775,9 @@ def test_cache_location(self):
         self.assertTrue(hasattr(py_zoneinfo.ZoneInfo, "_weak_cache"))
 
     def test_gc_tracked(self):
-        # The pure Python version is tracked by the GC but (for now) the C
-        # version is not.
         import gc
-
         self.assertTrue(gc.is_tracked(py_zoneinfo.ZoneInfo))
-        self.assertFalse(gc.is_tracked(c_zoneinfo.ZoneInfo))
+        self.assertTrue(gc.is_tracked(c_zoneinfo.ZoneInfo))
 
 
 @dataclasses.dataclass(frozen=True)

diff --git a/Modules/_zoneinfo.c b/Modules/_zoneinfo.c
@@ -79,7 +79,7 @@ struct StrongCacheNode {
     PyObject *zone;
 };
 
-static PyTypeObject PyZoneInfo_ZoneInfoType;
+static PyTypeObject *PyZoneInfo_ZoneInfoType;
 
 // Globals
 static PyObject *TIMEDELTA_CACHE = NULL;
@@ -242,7 +242,7 @@ zoneinfo_new_instance(PyTypeObject *type, PyObject *key)
 static PyObject *
 get_weak_cache(PyTypeObject *type)
 {
-    if (type == &PyZoneInfo_ZoneInfoType) {
+    if (type == PyZoneInfo_ZoneInfoType) {
         return ZONEINFO_WEAK_CACHE;
     }
     else {
@@ -296,10 +296,19 @@ zoneinfo_new(PyTypeObject *type, PyObject *args, PyObject *kw)
     return instance;
 }
 
+static int
+zoneinfo_traverse(PyObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(Py_TYPE(self));
+    return 0;
+}
+
 static void
 zoneinfo_dealloc(PyObject *obj_self)
 {
+    PyObject_GC_UnTrack(obj_self);
     PyZoneInfo_ZoneInfo *self = (PyZoneInfo_ZoneInfo *)obj_self;
+    PyTypeObject *tp = Py_TYPE(obj_self);
 
     if (self->weakreflist != NULL) {
         PyObject_ClearWeakRefs(obj_self);
@@ -331,7 +340,8 @@ zoneinfo_dealloc(PyObject *obj_self)
     Py_XDECREF(self->key);
     Py_XDECREF(self->file_repr);
 
-    Py_TYPE(self)->tp_free((PyObject *)self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
 }
 
 static PyObject *
@@ -2366,7 +2376,7 @@ find_in_strong_cache(const StrongCacheNode *const root, PyObject *const key)
 static int
 eject_from_strong_cache(const PyTypeObject *const type, PyObject *key)
 {
-    if (type != &PyZoneInfo_ZoneInfoType) {
+    if (type != PyZoneInfo_ZoneInfoType) {
         return 0;
     }
 
@@ -2419,7 +2429,7 @@ move_strong_cache_node_to_front(StrongCacheNode **root, StrongCacheNode *node)
 static PyObject *
 zone_from_strong_cache(const PyTypeObject *const type, PyObject *const key)
 {
-    if (type != &PyZoneInfo_ZoneInfoType) {
+    if (type != PyZoneInfo_ZoneInfoType) {
         return NULL;  // Strong cache currently only implemented for base class
     }
 
@@ -2444,7 +2454,7 @@ static void
 update_strong_cache(const PyTypeObject *const type, PyObject *key,
                     PyObject *zone)
 {
-    if (type != &PyZoneInfo_ZoneInfoType) {
+    if (type != PyZoneInfo_ZoneInfoType) {
         return;
     }
 
@@ -2477,7 +2487,7 @@ update_strong_cache(const PyTypeObject *const type, PyObject *key,
 void
 clear_strong_cache(const PyTypeObject *const type)
 {
-    if (type != &PyZoneInfo_ZoneInfoType) {
+    if (type != PyZoneInfo_ZoneInfoType) {
         return;
     }
 
@@ -2585,23 +2595,31 @@ static PyMemberDef zoneinfo_members[] = {
      .type = T_OBJECT_EX,
      .flags = READONLY,
      .doc = NULL},
+    {.name = "__weaklistoffset__",
+     .offset = offsetof(PyZoneInfo_ZoneInfo, weakreflist),
+     .type = T_PYSSIZET,
+     .flags = READONLY},
     {NULL}, /* Sentinel */
 };
 
-static PyTypeObject PyZoneInfo_ZoneInfoType = {
-    PyVarObject_HEAD_INIT(NULL, 0)  //
-        .tp_name = "zoneinfo.ZoneInfo",
-    .tp_basicsize = sizeof(PyZoneInfo_ZoneInfo),
-    .tp_weaklistoffset = offsetof(PyZoneInfo_ZoneInfo, weakreflist),
-    .tp_repr = (reprfunc)zoneinfo_repr,
-    .tp_str = (reprfunc)zoneinfo_str,
-    .tp_getattro = PyObject_GenericGetAttr,
-    .tp_flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE),
-    /* .tp_doc = zoneinfo_doc, */
-    .tp_methods = zoneinfo_methods,
-    .tp_members = zoneinfo_members,
-    .tp_new = zoneinfo_new,
-    .tp_dealloc = zoneinfo_dealloc,
+static PyType_Slot zoneinfo_slots[] = {
+    {Py_tp_repr, zoneinfo_repr},
+    {Py_tp_str, zoneinfo_str},
+    {Py_tp_getattro, PyObject_GenericGetAttr},
+    {Py_tp_methods, zoneinfo_methods},
+    {Py_tp_members, zoneinfo_members},
+    {Py_tp_new, zoneinfo_new},
+    {Py_tp_dealloc, zoneinfo_dealloc},
+    {Py_tp_traverse, zoneinfo_traverse},
+    {0, NULL},
+};
+
+static PyType_Spec zoneinfo_spec = {
+    .name = "zoneinfo.ZoneInfo",
+    .basicsize = sizeof(PyZoneInfo_ZoneInfo),
+    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC),
+    .slots = zoneinfo_slots,
+    /* .doc = zoneinfo_doc, */
 };
 
 /////
@@ -2633,7 +2651,7 @@ module_free()
         Py_CLEAR(ZONEINFO_WEAK_CACHE);
     }
 
-    clear_strong_cache(&PyZoneInfo_ZoneInfoType);
+    clear_strong_cache(PyZoneInfo_ZoneInfoType);
 }
 
 static int
@@ -2643,13 +2661,19 @@ zoneinfomodule_exec(PyObject *m)
     if (PyDateTimeAPI == NULL) {
         goto error;
     }
-    PyZoneInfo_ZoneInfoType.tp_base = PyDateTimeAPI->TZInfoType;
-    if (PyType_Ready(&PyZoneInfo_ZoneInfoType) < 0) {
+    PyObject *base = (PyObject *)PyDateTimeAPI->TZInfoType;
+    PyZoneInfo_ZoneInfoType = (PyTypeObject *)PyType_FromModuleAndSpec(m,
+                                                                       &zoneinfo_spec,
+                                                                       base);
+    if (PyZoneInfo_ZoneInfoType == NULL) {
         goto error;
     }
 
-    Py_INCREF(&PyZoneInfo_ZoneInfoType);
-    PyModule_AddObject(m, "ZoneInfo", (PyObject *)&PyZoneInfo_ZoneInfoType);
+    int rc = PyModule_AddObjectRef(m, "ZoneInfo", (PyObject *)PyZoneInfo_ZoneInfoType);
+    if (rc < 0) {
+        goto error;
+    }
+    Py_DECREF(PyZoneInfo_ZoneInfoType);
 
     /* Populate imports */
     PyObject *_tzpath_module = PyImport_ImportModule("zoneinfo._tzpath");