traverseda
/
cpython
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Issue 5032: added a step argument to itertools.count() and allowed non-integer arguments.

This commit is contained in:
Raymond Hettinger 2009-02-12 06:28:27 +00:00
parent a9cab51965
commit 3072921d0e
4 changed files with 133 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
Doc/library/itertools.rst
View File

@ -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)

37
Lib/test/test_itertools.py
View File

@ -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('')), [])

3
Misc/NEWS
View File

@ -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.

114
Modules/itertoolsmodule.c
View File

@ -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)
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;
}
if (one == NULL) {
one = PyLong_FromLong(1);
if (one == 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)