From 3662c9090db9831d689de006091d59b069849592 Mon Sep 17 00:00:00 2001
From: Raymond Hettinger <python@rcn.com>
Date: Thu, 6 Mar 2008 22:58:42 +0000
Subject: [PATCH] Backport r61286 adding GC to the grouper for
 itertools.groupby() fixing Issue 2246.

---
 Lib/test/test_itertools.py |  7 +++++++
 Modules/itertoolsmodule.c  | 20 +++++++++++++++-----
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 99f9077013c..e949cf6ce76 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -449,6 +449,13 @@ class TestGC(unittest.TestCase):
         a = []
         self.makecycle(groupby([a]*2, lambda x:x), a)
 
+    def test_issue2246(self):
+        # Issue 2246 -- the _grouper iterator was not included in GC
+        n = 10
+        keyfunc = lambda x: x
+        for i, j in groupby(xrange(n), key=keyfunc):
+            keyfunc.__dict__.setdefault('x',[]).append(j)
+
     def test_ifilter(self):
         a = []
         self.makecycle(ifilter(lambda x:True, [a]*2), a)
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 31f20c710c7..55731eb92f0 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -199,7 +199,7 @@ _grouper_create(groupbyobject *parent, PyObject *tgtkey)
 {
 	_grouperobject *igo;
 
-	igo = PyObject_New(_grouperobject, &_grouper_type);
+	igo = PyObject_GC_New(_grouperobject, &_grouper_type);
 	if (igo == NULL)
 		return NULL;
 	igo->parent = (PyObject *)parent;
@@ -207,15 +207,25 @@ _grouper_create(groupbyobject *parent, PyObject *tgtkey)
 	igo->tgtkey = tgtkey;
 	Py_INCREF(tgtkey);
 
+	PyObject_GC_Track(igo);
 	return (PyObject *)igo;
 }
 
 static void
 _grouper_dealloc(_grouperobject *igo)
 {
+	PyObject_GC_UnTrack(igo);
 	Py_DECREF(igo->parent);
 	Py_DECREF(igo->tgtkey);
-	PyObject_Del(igo);
+	PyObject_GC_Del(igo);
+}
+
+static int
+_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
+{
+	Py_VISIT(igo->parent);
+	Py_VISIT(igo->tgtkey);
+	return 0;
 }
 
 static PyObject *
@@ -282,9 +292,9 @@ static PyTypeObject _grouper_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 */
 	0,				/* tp_doc */
-	0, 				/* tp_traverse */
+	(traverseproc)_grouper_traverse,/* tp_traverse */
 	0,				/* tp_clear */
 	0,				/* tp_richcompare */
 	0,				/* tp_weaklistoffset */
@@ -301,7 +311,7 @@ static PyTypeObject _grouper_type = {
 	0,				/* tp_init */
 	0,				/* tp_alloc */
 	0,				/* tp_new */
-	PyObject_Del,			/* tp_free */
+	PyObject_GC_Del,		/* tp_free */
 };