Commit 6d121f16 authored by Raymond Hettinger's avatar Raymond Hettinger

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

parent de33c624
...@@ -52,8 +52,7 @@ class TestBasicOps(unittest.TestCase): ...@@ -52,8 +52,7 @@ class TestBasicOps(unittest.TestCase):
self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)]) self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
self.assertRaises(TypeError, count, 2, 3) self.assertRaises(TypeError, count, 2, 3)
self.assertRaises(TypeError, count, 'a') self.assertRaises(TypeError, count, 'a')
c = count(sys.maxint-2) # verify that rollover doesn't crash self.assertRaises(OverflowError, list, islice(count(sys.maxint-5), 10))
c.next(); c.next(); c.next(); c.next(); c.next()
c = count(3) c = count(3)
self.assertEqual(repr(c), 'count(3)') self.assertEqual(repr(c), 'count(3)')
c.next() c.next()
......
...@@ -2073,6 +2073,11 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) ...@@ -2073,6 +2073,11 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
static PyObject * static PyObject *
count_next(countobject *lz) 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++); return PyInt_FromSsize_t(lz->cnt++);
} }
......
...@@ -62,6 +62,12 @@ enum_next(enumobject *en) ...@@ -62,6 +62,12 @@ enum_next(enumobject *en)
PyObject *result = en->en_result; PyObject *result = en->en_result;
PyObject *it = en->en_sit; 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); next_item = (*it->ob_type->tp_iternext)(it);
if (next_item == NULL) if (next_item == NULL)
return NULL; return NULL;
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment