From 50e90e265fc3205fc6cc0193d3cebe08e9859b0a Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Thu, 4 Oct 2007 00:20:27 +0000 Subject: itertools.count() no longer limited to sys.maxint. --- Doc/library/itertools.rst | 10 ++---- Lib/test/test_itertools.py | 7 +++- Misc/NEWS | 3 ++ Modules/itertoolsmodule.c | 79 ++++++++++++++++++++++++++++++++++++++++------ 4 files changed, 82 insertions(+), 17 deletions(-) diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst index cb168a8..93e62f6 100644 --- a/Doc/library/itertools.rst +++ b/Doc/library/itertools.rst @@ -79,19 +79,15 @@ loops that truncate the stream. .. function:: count([n]) Make an iterator that returns consecutive integers starting with *n*. If not - specified *n* defaults to zero. Does not currently support python long - integers. Often used as an argument to :func:`imap` to generate consecutive - data points. Also, used with :func:`izip` to add sequence numbers. Equivalent - to:: + specified *n* defaults to zero. Often used as an argument to :func:`imap` to + generate consecutive data points. Also, used with :func:`izip` to add sequence + numbers. Equivalent to:: def count(n=0): while True: yield n n += 1 - Note, :func:`count` does not check for overflow and will return negative numbers - after exceeding ``sys.maxint``. This behavior may change in the future. - .. function:: cycle(iterable) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 6c362ad..3a370d9 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -52,9 +52,12 @@ class TestBasicOps(unittest.TestCase): self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)]) self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)]) self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)]) + self.assertEqual(take(2, zip('abc',count(-1))), [('a', -1), ('b', 0)]) + self.assertEqual(take(2, zip('abc',count(-3))), [('a', -3), ('b', -2)]) self.assertRaises(TypeError, count, 2, 3) self.assertRaises(TypeError, count, 'a') - self.assertRaises(OverflowError, list, islice(count(maxsize-5), 10)) + self.assertEqual(list(islice(count(maxsize-5), 10)), range(maxsize-5, maxsize+5)) + self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5)) c = count(3) self.assertEqual(repr(c), 'count(3)') c.next() @@ -63,6 +66,8 @@ class TestBasicOps(unittest.TestCase): self.assertEqual(repr(c), 'count(-9)') c.next() self.assertEqual(c.next(), -8) + for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5): + self.assertEqual(repr(count(i)), 'count(%r)' % i) def test_cycle(self): self.assertEqual(take(10, cycle('abc')), list('abcabcabca')) diff --git a/Misc/NEWS b/Misc/NEWS index 01f860b..bd31af4 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -270,6 +270,9 @@ Core and builtins Library ------- +- itertools.count() is no longer bounded to LONG_MAX. Formerly, it raised + an OverflowError. Now, automatically shifts from ints to longs. + - Patch #1541463: optimize performance of cgi.FieldStorage operations. - Decimal is fully updated to the latest Decimal Specification (v1.66). diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 35b77e7..fb54b14 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -2032,6 +2032,7 @@ static PyTypeObject ifilterfalse_type = { typedef struct { PyObject_HEAD Py_ssize_t cnt; + PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */ } countobject; static PyTypeObject count_type; @@ -2041,37 +2042,97 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { countobject *lz; Py_ssize_t cnt = 0; + PyObject *cnt_arg = NULL; + PyObject *long_cnt = NULL; if (type == &count_type && !_PyArg_NoKeywords("count()", kwds)) return NULL; - if (!PyArg_ParseTuple(args, "|n:count", &cnt)) + if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg)) return NULL; + if (cnt_arg != NULL) { + cnt = PyInt_AsSsize_t(cnt_arg); + if (cnt == -1 && PyErr_Occurred()) { + PyErr_Clear(); + if (!PyLong_Check(cnt_arg)) { + PyErr_SetString(PyExc_TypeError, "an integer is required"); + return NULL; + } + long_cnt = cnt_arg; + Py_INCREF(long_cnt); + cnt = PY_SSIZE_T_MAX; + } + } + /* create countobject structure */ lz = (countobject *)PyObject_New(countobject, &count_type); - if (lz == NULL) + if (lz == NULL) { + Py_XDECREF(long_cnt); return NULL; + } lz->cnt = cnt; + lz->long_cnt = long_cnt; return (PyObject *)lz; } +static void +count_dealloc(countobject *lz) +{ + Py_XDECREF(lz->long_cnt); + PyObject_Del(lz); +} + +static PyObject * +count_nextlong(countobject *lz) +{ + static PyObject *one = NULL; + PyObject *cnt; + PyObject *stepped_up; + + if (lz->long_cnt == NULL) { + lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX); + if (lz->long_cnt == NULL) + return NULL; + } + if (one == NULL) { + one = PyInt_FromLong(1); + if (one == NULL) + return NULL; + } + cnt = lz->long_cnt; + assert(cnt != NULL); + stepped_up = PyNumber_Add(cnt, one); + if (stepped_up == NULL) + return NULL; + lz->long_cnt = stepped_up; + return cnt; +} + static PyObject * count_next(countobject *lz) { - if (lz->cnt == PY_SSIZE_T_MAX) { - PyErr_SetString(PyExc_OverflowError, - "cannot count beyond PY_SSIZE_T_MAX"); - return NULL; - } + if (lz->cnt == PY_SSIZE_T_MAX) + return count_nextlong(lz); return PyInt_FromSsize_t(lz->cnt++); } static PyObject * count_repr(countobject *lz) { - return PyString_FromFormat("count(%zd)", lz->cnt); + PyObject *cnt_repr; + PyObject *result; + + if (lz->cnt != PY_SSIZE_T_MAX) + return PyString_FromFormat("count(%zd)", lz->cnt); + + cnt_repr = PyObject_Repr(lz->long_cnt); + if (cnt_repr == NULL) + return NULL; + result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr)); + Py_DECREF(cnt_repr); + return result; } PyDoc_STRVAR(count_doc, @@ -2086,7 +2147,7 @@ static PyTypeObject count_type = { sizeof(countobject), /* tp_basicsize */ 0, /* tp_itemsize */ /* methods */ - (destructor)PyObject_Del, /* tp_dealloc */ + (destructor)count_dealloc, /* tp_dealloc */ 0, /* tp_print */ 0, /* tp_getattr */ 0, /* tp_setattr */ -- cgit v0.12