diff options
| -rw-r--r-- | Doc/library/itertools.rst | 20 | ||||
| -rw-r--r-- | Lib/test/test_itertools.py | 37 | ||||
| -rw-r--r-- | Misc/NEWS | 3 | ||||
| -rw-r--r-- | Modules/itertoolsmodule.c | 114 |
4 files changed, 133 insertions, 41 deletions
diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst index 8346587..fda3beb 100644 --- a/Doc/library/itertools.rst +++ b/Doc/library/itertools.rst @@ -178,7 +178,7 @@ loops that truncate the stream. The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``. - .. versionadded:: 2.7 + .. versionadded:: 3.1 .. function:: compress(data, selectors) @@ -191,22 +191,24 @@ loops that truncate the stream. # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F return (d for d, s in zip(data, selectors) if s) - .. versionadded:: 2.7 + .. versionadded:: 3.1 -.. function:: count([n]) +.. function:: count(n=0, step=1) - Make an iterator that returns consecutive integers starting with *n*. If not - specified *n* defaults to zero. Often used as an argument to :func:`map` to - generate consecutive data points. Also, used with :func:`zip` to add sequence - numbers. Equivalent to:: + Make an iterator that returns evenly spaced values starting with *n*. Often + used as an argument to :func:`map` to generate consecutive data points. + Also, used with :func:`zip` to add sequence numbers. Equivalent to:: - def count(n=0): + def count(n=0, step=1): # count(10) --> 10 11 12 13 14 ... + # count(2.5, 0.5) -> 3.5 3.0 4.5 ... while True: yield n - n += 1 + n += step + .. versionchanged:: 3.1 + added *step* argument and allowed non-integer arguments. .. function:: cycle(iterable) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 50c3402..a96d994 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -328,7 +328,7 @@ class TestBasicOps(unittest.TestCase): self.assertEqual(take(2, lzip('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, 2, 3, 4) self.assertRaises(TypeError, count, 'a') self.assertEqual(list(islice(count(maxsize-5), 10)), list(range(maxsize-5, maxsize+5))) @@ -341,6 +341,7 @@ class TestBasicOps(unittest.TestCase): c = count(-9) self.assertEqual(repr(c), 'count(-9)') next(c) + self.assertEqual(repr(count(10.25)), 'count(10.25)') self.assertEqual(next(c), -8) for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5): # Test repr (ignoring the L in longs) @@ -348,6 +349,40 @@ class TestBasicOps(unittest.TestCase): r2 = 'count(%r)'.__mod__(i).replace('L', '') self.assertEqual(r1, r2) + def test_count_with_stride(self): + self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)]) + self.assertEqual(lzip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)]) + self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)]) + self.assertEqual(take(20, count(maxsize-15, 3)), take(20, range(maxsize-15, maxsize+100, 3))) + self.assertEqual(take(20, count(-maxsize-15, 3)), take(20, range(-maxsize-15,-maxsize+100, 3))) + self.assertEqual(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j]) + self.assertEqual(repr(take(3, count(10, 2.5))), repr([10, 12.5, 15.0])) + c = count(3, 5) + self.assertEqual(repr(c), 'count(3, 5)') + next(c) + self.assertEqual(repr(c), 'count(8, 5)') + c = count(-9, 0) + self.assertEqual(repr(c), 'count(-9, 0)') + next(c) + self.assertEqual(repr(c), 'count(-9, 0)') + c = count(-9, -3) + self.assertEqual(repr(c), 'count(-9, -3)') + next(c) + self.assertEqual(repr(c), 'count(-12, -3)') + self.assertEqual(repr(c), 'count(-12, -3)') + self.assertEqual(repr(count(10.5, 1.25)), 'count(10.5, 1.25)') + self.assertEqual(repr(count(10.5, 1)), 'count(10.5)') # suppress step=1 when it's an int + self.assertEqual(repr(count(10.5, 1.00)), 'count(10.5, 1.0)') # do show float values lilke 1.0 + for i in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 10, sys.maxsize-5, sys.maxsize+5): + for j in (-sys.maxsize-5, -sys.maxsize+5 ,-10, -1, 0, 1, 10, sys.maxsize-5, sys.maxsize+5): + # Test repr (ignoring the L in longs) + r1 = repr(count(i, j)).replace('L', '') + if j == 1: + r2 = ('count(%r)' % i).replace('L', '') + else: + r2 = ('count(%r, %r)' % (i, j)).replace('L', '') + self.assertEqual(r1, r2) + def test_cycle(self): self.assertEqual(take(10, cycle('abc')), list('abcabcabca')) self.assertEqual(list(cycle('')), []) @@ -241,6 +241,9 @@ Library - Added a new itertools functions: combinations_with_replacement() and compress(). +- Issue 5032: added a step argument to itertools.count() and + allowed non-integer arguments. + - Fix and properly document the multiprocessing module's logging support, expose the internal levels and provide proper usage examples. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 9245b1f..284b960 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -2886,9 +2886,27 @@ static PyTypeObject filterfalse_type = { typedef struct { PyObject_HEAD Py_ssize_t cnt; - PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */ + PyObject *long_cnt; + PyObject *long_step; } countobject; +/* Counting logic and invariants: + +C_add_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified. + + assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1)); + Advances with: cnt += 1 + When count hits Y_SSIZE_T_MAX, switch to Py_add_mode. + +Py_add_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float. + + assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL); + All counting is done with python objects (no overflows or underflows). + Advances with: long_cnt += long_step + Step may be zero -- effectively a slow version of repeat(cnt). + Either long_cnt or long_step may be a float. +*/ + static PyTypeObject count_type; static PyObject * @@ -2896,28 +2914,45 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { countobject *lz; Py_ssize_t cnt = 0; - PyObject *cnt_arg = NULL; PyObject *long_cnt = NULL; + PyObject *long_step = NULL; if (type == &count_type && !_PyArg_NoKeywords("count()", kwds)) return NULL; - if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg)) + if (!PyArg_UnpackTuple(args, "count", 0, 2, &long_cnt, &long_step)) return NULL; - if (cnt_arg != NULL) { - cnt = PyLong_AsSsize_t(cnt_arg); - if (cnt == -1 && PyErr_Occurred()) { + if (long_cnt != NULL && !PyNumber_Check(long_cnt) || + long_step != NULL && !PyNumber_Check(long_step)) { + PyErr_SetString(PyExc_TypeError, "a number is required"); + return NULL; + } + + if (long_step == NULL) { + /* If not specified, step defaults to 1 */ + long_step = PyLong_FromLong(1); + if (long_step == NULL) + return NULL; + } else + Py_INCREF(long_step); + assert(long_step != NULL); + + if (long_cnt != NULL) { + cnt = PyLong_AsSsize_t(long_cnt); + if ((cnt == -1 && PyErr_Occurred()) || + !PyIndex_Check(long_cnt) || + !PyLong_Check(long_step) || + PyLong_AS_LONG(long_step) != 1) { + /* Switch to Py_add_mode */ 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; - } + } else + long_cnt = NULL; } + assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL || + cnt == PY_SSIZE_T_MAX && long_cnt != NULL); /* create countobject structure */ lz = (countobject *)PyObject_New(countobject, &count_type); @@ -2927,6 +2962,7 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) } lz->cnt = cnt; lz->long_cnt = long_cnt; + lz->long_step = long_step; return (PyObject *)lz; } @@ -2934,7 +2970,8 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) static void count_dealloc(countobject *lz) { - Py_XDECREF(lz->long_cnt); + Py_XDECREF(lz->long_cnt); + Py_XDECREF(lz->long_step); PyObject_Del(lz); } @@ -2942,32 +2979,29 @@ static PyObject * count_nextlong(countobject *lz) { static PyObject *one = NULL; - PyObject *cnt; + PyObject *long_cnt; PyObject *stepped_up; - if (lz->long_cnt == NULL) { - lz->long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX); - if (lz->long_cnt == NULL) - return NULL; - } - if (one == NULL) { - one = PyLong_FromLong(1); - if (one == NULL) + long_cnt = lz->long_cnt; + if (long_cnt == NULL) { + /* Switch to Py_add_mode */ + long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX); + if (long_cnt == NULL) return NULL; } - cnt = lz->long_cnt; - assert(cnt != NULL); - stepped_up = PyNumber_Add(cnt, one); + assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL); + + stepped_up = PyNumber_Add(long_cnt, lz->long_step); if (stepped_up == NULL) return NULL; lz->long_cnt = stepped_up; - return cnt; + return long_cnt; } static PyObject * count_next(countobject *lz) { - if (lz->cnt == PY_SSIZE_T_MAX) + if (lz->cnt == PY_SSIZE_T_MAX) return count_nextlong(lz); return PyLong_FromSsize_t(lz->cnt++); } @@ -2975,17 +3009,35 @@ count_next(countobject *lz) static PyObject * count_repr(countobject *lz) { - if (lz->cnt != PY_SSIZE_T_MAX) + if (lz->cnt != PY_SSIZE_T_MAX) return PyUnicode_FromFormat("count(%zd)", lz->cnt); - return PyUnicode_FromFormat("count(%R)", lz->long_cnt); + if (PyLong_Check(lz->long_step)) { + long step = PyLong_AsLong(lz->long_step); + if (step == -1 && PyErr_Occurred()) { + PyErr_Clear(); + } + if (step == 1) { + /* Don't display step when it is an integer equal to 1 */ + return PyUnicode_FromFormat("count(%R)", lz->long_cnt); + } + } + return PyUnicode_FromFormat("count(%R, %R)", + lz->long_cnt, lz->long_step); } PyDoc_STRVAR(count_doc, -"count([firstval]) --> count object\n\ + "count([firstval[, step]]) --> count object\n\ \n\ Return a count object whose .__next__() method returns consecutive\n\ -integers starting from zero or, if specified, from firstval."); +integers starting from zero or, if specified, from firstval.\n\ +If step is specified, counts by that interval.\n\ +Same as:\n\ + def count(firstval=0, step=1):\n\ + x = firstval\n\ + while 1:\n\ + yield x\n\ + x += step\n"); static PyTypeObject count_type = { PyVarObject_HEAD_INIT(NULL, 0) |