Make itertools.tee and its internal teedataobject participate in GC. This

alone does not solve the leak in test_generators, unfortunately, but it is
part of test_generators' problem and it does solve other cycles.

1 files changed, 71 insertions, 19 deletions

diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 49d241f..71081fb 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -340,7 +340,7 @@ teedataobject_new(PyObject *it)
 {
 	teedataobject *tdo;
 
-	tdo = PyObject_New(teedataobject, &teedataobject_type);
+	tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
 	if (tdo == NULL)
 		return NULL;
 
@@ -348,6 +348,7 @@ teedataobject_new(PyObject *it)
 	tdo->nextlink = NULL;
 	Py_INCREF(it);
 	tdo->it = it;
+	PyObject_GC_Track(tdo);
 	return (PyObject *)tdo;
 }
 
@@ -381,16 +382,34 @@ teedataobject_getitem(teedataobject *tdo, int i)
 	return value;
 }
 
-static void
-teedataobject_dealloc(teedataobject *tdo)
+static int
+teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
 {
 	int i;
+	Py_VISIT(tdo->it);
+	for (i = 0; i < tdo->numread; i++)
+		Py_VISIT(tdo->values[i]);
+	Py_VISIT(tdo->nextlink);
+	return 0;
+}
 
+static int
+teedataobject_clear(teedataobject *tdo)
+{
+	int i;
+	Py_CLEAR(tdo->it);
 	for (i=0 ; i<tdo->numread ; i++)
-		Py_DECREF(tdo->values[i]);
-	Py_XDECREF(tdo->it);
-	Py_XDECREF(tdo->nextlink);
-	PyObject_Del(tdo);
+		Py_CLEAR(tdo->values[i]);
+	Py_CLEAR(tdo->nextlink);
+	return 0;
+}
+
+static void
+teedataobject_dealloc(teedataobject *tdo)
+{
+	PyObject_GC_UnTrack(tdo);
+	teedataobject_clear(tdo);
+	PyObject_GC_Del(tdo);
 }
 
 PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
@@ -417,9 +436,26 @@ static PyTypeObject teedataobject_type = {
 	PyObject_GenericGetAttr,		/* tp_getattro */
 	0,					/* tp_setattro */
 	0,					/* tp_as_buffer */
-	Py_TPFLAGS_DEFAULT,			/* tp_flags */
+	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,	/* tp_flags */
 	teedataobject_doc,			/* tp_doc */
-	0,					/* tp_traverse */
+	(traverseproc)teedataobject_traverse,	/* tp_traverse */
+	(inquiry)teedataobject_clear,		/* tp_clear */
+	0,					/* tp_richcompare */
+	0,					/* tp_weaklistoffset */
+	0,					/* tp_iter */
+	0,					/* tp_iternext */
+	0,					/* tp_methods */
+	0,					/* tp_members */
+	0,					/* tp_getset */
+	0,					/* tp_base */
+	0,					/* tp_dict */
+	0,					/* tp_descr_get */
+	0,					/* tp_descr_set */
+	0,					/* tp_dictoffset */
+	0,					/* tp_init */
+	0,					/* tp_alloc */
+	0,					/* tp_new */
+	PyObject_GC_Del,			/* tp_free */
 };
 
 
@@ -443,12 +479,19 @@ tee_next(teeobject *to)
 	return value;
 }
 
+static int
+tee_traverse(teeobject *to, visitproc visit, void *arg)
+{
+	Py_VISIT((PyObject *)to->dataobj);
+	return 0;
+}
+
 static PyObject *
 tee_copy(teeobject *to)
 {
 	teeobject *newto;
 
-	newto = PyObject_New(teeobject, &tee_type);
+	newto = PyObject_GC_New(teeobject, &tee_type);
 	if (newto == NULL)
 		return NULL;
 	Py_INCREF(to->dataobj);
@@ -474,12 +517,13 @@ tee_fromiterable(PyObject *iterable)
 		goto done;
 	}
 
-	to = PyObject_New(teeobject, &tee_type);
+	to = PyObject_GC_New(teeobject, &tee_type);
 	if (to == NULL) 
 		goto done;
 	to->dataobj = (teedataobject *)teedataobject_new(it);
 	to->index = 0;
 	to->weakreflist = NULL;
+	PyObject_GC_Track(to);
 done:
 	Py_XDECREF(it);
 	return (PyObject *)to;
@@ -495,13 +539,21 @@ tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 	return tee_fromiterable(iterable);
 }
 
-static void
-tee_dealloc(teeobject *to)
+static int
+tee_clear(teeobject *to)
 {
 	if (to->weakreflist != NULL)
 		PyObject_ClearWeakRefs((PyObject *) to);
-	Py_XDECREF(to->dataobj);
-	PyObject_Del(to);
+	Py_CLEAR(to->dataobj);
+	return 0;
+}
+
+static void
+tee_dealloc(teeobject *to)
+{
+	PyObject_GC_UnTrack(to);
+	tee_clear(to);
+	PyObject_GC_Del(to);
 }
 
 PyDoc_STRVAR(teeobject_doc,
@@ -534,10 +586,10 @@ static PyTypeObject tee_type = {
 	0,				/* tp_getattro */
 	0,				/* tp_setattro */
 	0,				/* tp_as_buffer */
-	Py_TPFLAGS_DEFAULT,		/* tp_flags */
+	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,	/* tp_flags */
 	teeobject_doc,			/* tp_doc */
-	0,				/* tp_traverse */
-	0,				/* tp_clear */
+	(traverseproc)tee_traverse,	/* tp_traverse */
+	(inquiry)tee_clear,		/* tp_clear */
 	0,				/* tp_richcompare */
 	offsetof(teeobject, weakreflist),	/* tp_weaklistoffset */
 	PyObject_SelfIter,		/* tp_iter */
@@ -553,7 +605,7 @@ static PyTypeObject tee_type = {
 	0,				/* tp_init */
 	0,				/* tp_alloc */
 	tee_new,			/* tp_new */
-	PyObject_Del,			/* tp_free */
+	PyObject_GC_Del,		/* tp_free */
 };
 
 static PyObject *