Add GC support to count() objects.

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 812a1efdb48089aafc13349270b7eef15395c4ff..c34b915cebf98e4faa80667d8a90afbe3c38a8cd 100644 (file)
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -979,6 +979,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 a44d2ae3db9b120bd1913fb03bef616ce226b677..3cec4212918111eec74191786add4db113dff84b 100644 (file)
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -2954,7 +2954,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;
@@ -2969,9 +2969,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 *
@@ -3058,9 +3066,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 */
@@ -3077,6 +3086,7 @@ static PyTypeObject count_type = {
        0,                              /* tp_init */
        0,                              /* tp_alloc */
        count_new,                      /* tp_new */
+       PyObject_GC_Del,                /* tp_free */
 };