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

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('')), [])
diff --git a/Misc/NEWS b/Misc/NEWS
index 1be4001..7f99d45 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -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)