bpo-42972: Fully implement GC protocol for functools keywrapper and partial types (GH-26363)

1 files changed, 37 insertions, 22 deletions

diff --git a/Modules/_functoolsmodule.c b/Modules/_functoolsmodule.c
index 1cfb08c..b028f8f 100644
--- a/Modules/_functoolsmodule.c
+++ b/Modules/_functoolsmodule.c
@@ -147,18 +147,37 @@ partial_new(PyTypeObject *type, PyObject *args, PyObject *kw)
     return (PyObject *)pto;
 }
 
+static int
+partial_clear(partialobject *pto)
+{
+    Py_CLEAR(pto->fn);
+    Py_CLEAR(pto->args);
+    Py_CLEAR(pto->kw);
+    Py_CLEAR(pto->dict);
+    return 0;
+}
+
+static int
+partial_traverse(partialobject *pto, visitproc visit, void *arg)
+{
+    Py_VISIT(Py_TYPE(pto));
+    Py_VISIT(pto->fn);
+    Py_VISIT(pto->args);
+    Py_VISIT(pto->kw);
+    Py_VISIT(pto->dict);
+    return 0;
+}
+
 static void
 partial_dealloc(partialobject *pto)
 {
     PyTypeObject *tp = Py_TYPE(pto);
     /* bpo-31095: UnTrack is needed before calling any callbacks */
     PyObject_GC_UnTrack(pto);
-    if (pto->weakreflist != NULL)
+    if (pto->weakreflist != NULL) {
         PyObject_ClearWeakRefs((PyObject *) pto);
-    Py_XDECREF(pto->fn);
-    Py_XDECREF(pto->args);
-    Py_XDECREF(pto->kw);
-    Py_XDECREF(pto->dict);
+    }
+    (void)partial_clear(pto);
     tp->tp_free(pto);
     Py_DECREF(tp);
 }
@@ -307,16 +326,6 @@ partial_call(partialobject *pto, PyObject *args, PyObject *kwargs)
     return res;
 }
 
-static int
-partial_traverse(partialobject *pto, visitproc visit, void *arg)
-{
-    Py_VISIT(pto->fn);
-    Py_VISIT(pto->args);
-    Py_VISIT(pto->kw);
-    Py_VISIT(pto->dict);
-    return 0;
-}
-
 PyDoc_STRVAR(partial_doc,
 "partial(func, *args, **keywords) - new function with partial application\n\
     of the given arguments and keywords.\n");
@@ -469,6 +478,7 @@ static PyType_Slot partial_type_slots[] = {
     {Py_tp_setattro, PyObject_GenericSetAttr},
     {Py_tp_doc, (void *)partial_doc},
     {Py_tp_traverse, partial_traverse},
+    {Py_tp_clear, partial_clear},
     {Py_tp_methods, partial_methods},
     {Py_tp_members, partial_memberlist},
     {Py_tp_getset, partial_getsetlist},
@@ -506,14 +516,16 @@ static void
 keyobject_dealloc(keyobject *ko)
 {
     PyTypeObject *tp = Py_TYPE(ko);
-    keyobject_clear(ko);
-    PyObject_Free(ko);
+    PyObject_GC_UnTrack(ko);
+    (void)keyobject_clear(ko);
+    tp->tp_free(ko);
     Py_DECREF(tp);
 }
 
 static int
 keyobject_traverse(keyobject *ko, visitproc visit, void *arg)
 {
+    Py_VISIT(Py_TYPE(ko));
     Py_VISIT(ko->cmp);
     Py_VISIT(ko->object);
     return 0;
@@ -546,7 +558,8 @@ static PyType_Slot keyobject_type_slots[] = {
 static PyType_Spec keyobject_type_spec = {
     .name = "functools.KeyWrapper",
     .basicsize = sizeof(keyobject),
-    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION,
+    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION |
+              Py_TPFLAGS_HAVE_GC),
     .slots = keyobject_type_slots
 };
 
@@ -560,7 +573,7 @@ keyobject_call(keyobject *ko, PyObject *args, PyObject *kwds)
     if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:K", kwargs, &object))
         return NULL;
 
-    result = PyObject_New(keyobject, Py_TYPE(ko));
+    result = PyObject_GC_New(keyobject, Py_TYPE(ko));
     if (result == NULL) {
         return NULL;
     }
@@ -568,6 +581,7 @@ keyobject_call(keyobject *ko, PyObject *args, PyObject *kwds)
     result->cmp = ko->cmp;
     Py_INCREF(object);
     result->object = object;
+    PyObject_GC_Track(result);
     return (PyObject *)result;
 }
 
@@ -621,12 +635,13 @@ functools_cmp_to_key(PyObject *self, PyObject *args, PyObject *kwds)
         return NULL;
 
     state = get_functools_state(self);
-    object = PyObject_New(keyobject, state->keyobject_type);
+    object = PyObject_GC_New(keyobject, state->keyobject_type);
     if (!object)
         return NULL;
     Py_INCREF(cmp);
     object->cmp = cmp;
     object->object = NULL;
+    PyObject_GC_Track(object);
     return (PyObject *)object;
 }
 
@@ -754,7 +769,7 @@ lru_list_elem_dealloc(lru_list_elem *link)
     PyTypeObject *tp = Py_TYPE(link);
     Py_XDECREF(link->key);
     Py_XDECREF(link->result);
-    PyObject_Free(link);
+    tp->tp_free(link);
     Py_DECREF(tp);
 }
 
@@ -1260,7 +1275,7 @@ lru_cache_dealloc(lru_cache_object *obj)
         PyObject_ClearWeakRefs((PyObject*)obj);
     }
 
-    lru_cache_tp_clear(obj);
+    (void)lru_cache_tp_clear(obj);
     tp->tp_free(obj);
     Py_DECREF(tp);
 }