[Python-checkins] r69685 - in python/trunk: Lib/test/test_itertools.py Modules/itertoolsmodule.c

[raymond.hettinger]

Author: raymond.hettinger
Date: Mon Feb 16 21:39:12 2009
New Revision: 69685

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

Modified:
   python/trunk/Lib/test/test_itertools.py
   python/trunk/Modules/itertoolsmodule.c

Modified: python/trunk/Lib/test/test_itertools.py
==============================================================================
--- python/trunk/Lib/test/test_itertools.py	(original)
+++ python/trunk/Lib/test/test_itertools.py	Mon Feb 16 21:39:12 2009
@@ -966,6 +966,11 @@
         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)

Modified: python/trunk/Modules/itertoolsmodule.c
==============================================================================
--- python/trunk/Modules/itertoolsmodule.c	(original)
+++ python/trunk/Modules/itertoolsmodule.c	Mon Feb 16 21:39:12 2009
@@ -3271,7 +3271,7 @@
 		   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 @@
 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 @@
 	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 @@
 	0,				/* tp_init */
 	0,				/* tp_alloc */
 	count_new,			/* tp_new */
+	PyObject_GC_Del,		/* tp_free */
 };