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

Author: Erlend Egeberg Aasland

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,14 +573,15 @@ 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;
  }
  Py_INCREF(ko->cmp);
  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);
 }