diff options
| -rw-r--r-- | Doc/library/itertools.rst | 5 | ||||
| -rw-r--r-- | Lib/test/test_itertools.py | 1 | ||||
| -rw-r--r-- | Modules/itertoolsmodule.c | 11 |
3 files changed, 9 insertions, 8 deletions
diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst index fda3beb..ad5a23b 100644 --- a/Doc/library/itertools.rst +++ b/Doc/library/itertools.rst @@ -194,15 +194,16 @@ loops that truncate the stream. .. 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 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, step=1): + def count(start=0, step=1): # count(10) --> 10 11 12 13 14 ... # count(2.5, 0.5) -> 3.5 3.0 4.5 ... + n = start while True: yield n n += step diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 31ef2e1..812a1ef 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -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,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,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(3, count(2, 3.25-4j)), [2, 5.25-4j, 8.5-8j]) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 75c6e7e..a44d2ae 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -2916,11 +2916,10 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) Py_ssize_t cnt = 0; PyObject *long_cnt = NULL; PyObject *long_step = NULL; + static char *kwlist[] = {"start", "step", 0}; - if (type == &count_type && !_PyArg_NoKeywords("count()", kwds)) - return NULL; - - if (!PyArg_UnpackTuple(args, "count", 0, 2, &long_cnt, &long_step)) + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count", + kwlist, &long_cnt, &long_step)) return NULL; if (long_cnt != NULL && !PyNumber_Check(long_cnt) || @@ -3027,10 +3026,10 @@ count_repr(countobject *lz) } PyDoc_STRVAR(count_doc, - "count([firstval[, step]]) --> count object\n\ + "count([start[, step]]) --> count object\n\ \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\ def count(firstval=0, step=1):\n\ x = firstval\n\ |