Author: raymond.hettinger Date: Thu Oct 4 02:20:27 2007 New Revision: 58305 Modified: python/trunk/Doc/library/itertools.rst python/trunk/Lib/test/test_itertools.py python/trunk/Misc/NEWS python/trunk/Modules/itertoolsmodule.c Log: itertools.count() no longer limited to sys.maxint. Modified: python/trunk/Doc/library/itertools.rst ============================================================================== --- python/trunk/Doc/library/itertools.rst (original) +++ python/trunk/Doc/library/itertools.rst Thu Oct 4 02:20:27 2007 @@ -79,19 +79,15 @@ .. 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) Modified: python/trunk/Lib/test/test_itertools.py ============================================================================== --- python/trunk/Lib/test/test_itertools.py (original) +++ python/trunk/Lib/test/test_itertools.py Thu Oct 4 02:20:27 2007 @@ -52,9 +52,12 @@ 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 @@ 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')) Modified: python/trunk/Misc/NEWS ============================================================================== --- python/trunk/Misc/NEWS (original) +++ python/trunk/Misc/NEWS Thu Oct 4 02:20:27 2007 @@ -270,6 +270,9 @@ 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). Modified: python/trunk/Modules/itertoolsmodule.c ============================================================================== --- python/trunk/Modules/itertoolsmodule.c (original) +++ python/trunk/Modules/itertoolsmodule.c Thu Oct 4 02:20:27 2007 @@ -2032,6 +2032,7 @@ 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 @@ { 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 @@ 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 */