Do not let overflows in enumerate() and count() pass silently.

raymond.hettinger · raymond.hettinger · commit d48aab5d5bac · 2007-02-08T00:07:32.000Z
git-svn-id: http://svn.python.org/projects/python/trunk@53666 6015fed2-1504-0410-9fe1-9d1591cc4771
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
@@ -52,8 +52,7 @@ def test_count(self):
         self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
         self.assertRaises(TypeError, count, 2, 3)
         self.assertRaises(TypeError, count, 'a')
-        c = count(sys.maxint-2)   # verify that rollover doesn't crash
-        c.next(); c.next(); c.next(); c.next(); c.next()
+        self.assertRaises(OverflowError, list, islice(count(sys.maxint-5), 10))
         c = count(3)
         self.assertEqual(repr(c), 'count(3)')
         c.next()
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
@@ -2073,6 +2073,11 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 static PyObject *
 count_next(countobject *lz)
 {
+        if (lz->cnt == LONG_MAX) {
+                PyErr_SetString(PyExc_OverflowError,
+                        "cannot count beyond LONG_MAX");                
+                return NULL;         
+        }
 	return PyInt_FromSsize_t(lz->cnt++);
 }
 
diff --git a/Objects/enumobject.c b/Objects/enumobject.c
@@ -62,6 +62,12 @@ enum_next(enumobject *en)
 	PyObject *result = en->en_result;
 	PyObject *it = en->en_sit;
 
+	if (en->en_index == LONG_MAX) {
+		PyErr_SetString(PyExc_OverflowError,
+			"enumerate() is limited to LONG_MAX items");                
+		return NULL;         
+	}
+
 	next_item = (*it->ob_type->tp_iternext)(it);
 	if (next_item == NULL)
 		return NULL;

Original file line number	Diff line number	Diff line change
`@@ -2073,6 +2073,11 @@ count_new(PyTypeObject type, PyObject args, PyObject *kwds)`
`2073`	`2073`	`static PyObject *`
`2074`	`2074`	`count_next(countobject *lz)`
`2075`	`2075`	`{`
	`2076`	`+ if (lz->cnt == LONG_MAX) {`
	`2077`	`+ PyErr_SetString(PyExc_OverflowError,`
	`2078`	`+ "cannot count beyond LONG_MAX");`
	`2079`	`+ return NULL;`
	`2080`	`+ }`
`2076`	`2081`	`return PyInt_FromSsize_t(lz->cnt++);`
`2077`	`2082`	`}`
`2078`	`2083`