Add keyword argument support to itertools.count().

This commit is contained in:
Raymond Hettinger 2009-02-14 04:21:49 +00:00
parent 8c20189ce4
commit 9e8dbbcdcd
3 changed files with 9 additions and 8 deletions

View file

@ -194,15 +194,16 @@ loops that truncate the stream.
.. versionadded:: 3.1 .. versionadded:: 3.1
.. function:: count(n=0, step=1) .. function:: count(start=0, step=1)
Make an iterator that returns evenly spaced values starting with *n*. Often Make an iterator that returns evenly spaced values starting with *n*. Often
used as an argument to :func:`map` to generate consecutive data points. used as an argument to :func:`map` to generate consecutive data points.
Also, used with :func:`zip` to add sequence numbers. Equivalent to:: Also, used with :func:`zip` to add sequence numbers. Equivalent to::
def count(n=0, step=1): def count(start=0, step=1):
# count(10) --> 10 11 12 13 14 ... # count(10) --> 10 11 12 13 14 ...
# count(2.5, 0.5) -> 3.5 3.0 4.5 ... # count(2.5, 0.5) -> 3.5 3.0 4.5 ...
n = start
while True: while True:
yield n yield n
n += step n += step

View file

@ -355,6 +355,7 @@ class TestBasicOps(unittest.TestCase):
self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)]) 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,0)), [('a', 2), ('b', 2), ('c', 2)])
self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)]) self.assertEqual(lzip('abc',count(2,1)), [('a', 2), ('b', 3), ('c', 4)])
self.assertEqual(lzip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
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(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(take(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j])

View file

@ -2916,11 +2916,10 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Py_ssize_t cnt = 0; Py_ssize_t cnt = 0;
PyObject *long_cnt = NULL; PyObject *long_cnt = NULL;
PyObject *long_step = NULL; PyObject *long_step = NULL;
static char *kwlist[] = {"start", "step", 0};
if (type == &count_type && !_PyArg_NoKeywords("count()", kwds)) if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
return NULL; kwlist, &long_cnt, &long_step))
if (!PyArg_UnpackTuple(args, "count", 0, 2, &long_cnt, &long_step))
return NULL; return NULL;
if (long_cnt != NULL && !PyNumber_Check(long_cnt) || if (long_cnt != NULL && !PyNumber_Check(long_cnt) ||
@ -3027,10 +3026,10 @@ count_repr(countobject *lz)
} }
PyDoc_STRVAR(count_doc, PyDoc_STRVAR(count_doc,
"count([firstval[, step]]) --> count object\n\ "count([start[, step]]) --> count object\n\
\n\ \n\
Return a count object whose .__next__() method returns consecutive\n\ Return a count object whose .__next__() method returns consecutive\n\
integers starting from zero or, if specified, from firstval.\n\ integers starting from zero or, if specified, from start.\n\
If step is specified, counts by that interval. Equivalent to:\n\n\ If step is specified, counts by that interval. Equivalent to:\n\n\
def count(firstval=0, step=1):\n\ def count(firstval=0, step=1):\n\
x = firstval\n\ x = firstval\n\