Add GC support to count() objects. Backport candidate.

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index ce5c898788777e185000ffea5ca6ed0029e28172..a33939eb20f655ea99f56a4688bb44ae055deb93 100644 (file)
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -966,6 +966,11 @@ class TestGC(unittest.TestCase):
         a = []
         self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
 
+    def test_count(self):
+        a = []
+        Int = type('Int', (int,), dict(x=a))
+        self.makecycle(count(Int(0), Int(1)), a)
+
     def test_cycle(self):
         a = []
         self.makecycle(cycle([a]*2), a)

diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 8a40427f962163bcd1ae6245e2e5285b1ca8056d..7d1791b2e0bb78b1461cf103d38e57e925d512cc 100644 (file)
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -3271,7 +3271,7 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
                   cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
 
        /* create countobject structure */
-       lz = (countobject *)PyObject_New(countobject, &count_type);
+       lz = (countobject *)type->tp_alloc(type, 0);
        if (lz == NULL) {
                Py_XDECREF(long_cnt);
                return NULL;
@@ -3286,9 +3286,17 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 static void
 count_dealloc(countobject *lz)
 {
+       PyObject_GC_UnTrack(lz);
        Py_XDECREF(lz->long_cnt);
        Py_XDECREF(lz->long_step);
-       PyObject_Del(lz);
+       Py_TYPE(lz)->tp_free(lz);
+}
+
+count_traverse(countobject *lz, visitproc visit, void *arg)
+{
+       Py_VISIT(lz->long_cnt);
+       Py_VISIT(lz->long_step);
+       return 0;
 }
 
 static PyObject *
@@ -3384,9 +3392,10 @@ static PyTypeObject count_type = {
        PyObject_GenericGetAttr,        /* tp_getattro */
        0,                              /* tp_setattro */
        0,                              /* tp_as_buffer */
-       Py_TPFLAGS_DEFAULT,             /* tp_flags */
+       Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
+               Py_TPFLAGS_BASETYPE,            /* tp_flags */
        count_doc,                      /* tp_doc */
-       0,                              /* tp_traverse */
+       (traverseproc)count_traverse,                           /* tp_traverse */
        0,                              /* tp_clear */
        0,                              /* tp_richcompare */
        0,                              /* tp_weaklistoffset */
@@ -3403,6 +3412,7 @@ static PyTypeObject count_type = {
        0,                              /* tp_init */
        0,                              /* tp_alloc */
        count_new,                      /* tp_new */
+       PyObject_GC_Del,                /* tp_free */
 };