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

4 files changed, 135 insertions, 41 deletions

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index 240baf0..4dfc083 100644
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -200,20 +200,23 @@ loops that truncate the stream.
    .. versionadded:: 2.7
 
 
-.. 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:`imap` to
-   generate consecutive data points. Also, used with :func:`izip` to add sequence
-   numbers.  Equivalent to::
+   Make an iterator that returns evenly spaced values starting with *n*. Often
+   used as an argument to :func:`imap` to generate consecutive data points.
+   Also, used with :func:`izip` 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:: 2.7
+      added *step* argument and allowed non-integer arguments.
+
 .. function:: cycle(iterable)
 
    Make an iterator returning elements from the iterable and saving a copy of each.
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index b79f70f..9e6b7c8 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -324,7 +324,7 @@ class TestBasicOps(unittest.TestCase):
         self.assertEqual(take(2, zip('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)), range(maxsize-5, maxsize+5))
         self.assertEqual(list(islice(count(-maxsize-5), 10)), range(-maxsize-5, -maxsize+5))
@@ -335,6 +335,7 @@ class TestBasicOps(unittest.TestCase):
         c = count(-9)
         self.assertEqual(repr(c), 'count(-9)')
         c.next()
+        self.assertEqual(repr(count(10.25)), 'count(10.25)')
         self.assertEqual(c.next(), -8)
         for i in (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
             # Test repr (ignoring the L in longs)
@@ -342,6 +343,40 @@ class TestBasicOps(unittest.TestCase):
             r2 = 'count(%r)'.__mod__(i).replace('L', '')
             self.assertEqual(r1, r2)
 
+    def test_count_with_stride(self):
+        self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)])
+        self.assertEqual(zip('abc',count(2,0)), [('a', 2), ('b', 2), ('c', 2)])
+        self.assertEqual(zip('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)')
+        c.next()
+        self.assertEqual(repr(c), 'count(8, 5)')
+        c = count(-9, 0)
+        self.assertEqual(repr(c), 'count(-9, 0)')
+        c.next()
+        self.assertEqual(repr(c), 'count(-9, 0)')
+        c = count(-9, -3)
+        self.assertEqual(repr(c), 'count(-9, -3)')
+        c.next()
+        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.maxint-5, -sys.maxint+5 ,-10, -1, 0, 10, sys.maxint-5, sys.maxint+5):
+            for j in  (-sys.maxint-5, -sys.maxint+5 ,-10, -1, 0, 1, 10, sys.maxint-5, sys.maxint+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 530a55d..520edda 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -227,6 +227,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 486c2e0..aa764d8 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -3203,9 +3203,27 @@ static PyTypeObject ifilterfalse_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 *
@@ -3213,28 +3231,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 = PyInt_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 = PyInt_FromLong(1);
+		if (long_step == NULL)
+			return NULL;
+	} else
+		Py_INCREF(long_step);
+	assert(long_step != NULL);
+
+	if (long_cnt != NULL) {
+		cnt = PyInt_AsSsize_t(long_cnt);
+		if ((cnt == -1 && PyErr_Occurred()) || 
+				!PyIndex_Check(long_cnt)  || 
+				!PyInt_Check(long_step) ||
+				PyInt_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);
@@ -3244,6 +3279,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;
 }
@@ -3251,7 +3287,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);
 }
 
@@ -3259,32 +3296,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 = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
-		if (lz->long_cnt == NULL)
-			return NULL;
-	}
-	if (one == NULL) {
-		one = PyInt_FromLong(1);
-		if (one == NULL)
+	long_cnt = lz->long_cnt;
+	if (long_cnt == NULL) {
+		/* Switch to Py_add_mode */
+		long_cnt = PyInt_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 PyInt_FromSsize_t(lz->cnt++);
 }
@@ -3292,25 +3326,44 @@ count_next(countobject *lz)
 static PyObject *
 count_repr(countobject *lz)
 {
-	PyObject *cnt_repr;
-	PyObject *result;
+	PyObject *cnt_repr, *step_repr = NULL;
+	PyObject *result = NULL;
 
-        if (lz->cnt != PY_SSIZE_T_MAX)
+    if (lz->cnt != PY_SSIZE_T_MAX)
 		return PyString_FromFormat("count(%zd)", lz->cnt);
 
 	cnt_repr = PyObject_Repr(lz->long_cnt);
 	if (cnt_repr == NULL)
 		return NULL;
-	result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
+
+	if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
+			/* Don't display step when it is an integer equal to 1 */
+			result = PyString_FromFormat("count(%s)",
+										 PyString_AS_STRING(cnt_repr));
+	} else {
+		step_repr = PyObject_Repr(lz->long_step);
+		if (step_repr != NULL)
+			result = PyString_FromFormat("count(%s, %s)",
+										PyString_AS_STRING(cnt_repr),
+										PyString_AS_STRING(step_repr));
+	}
 	Py_DECREF(cnt_repr);
+	Py_XDECREF(step_repr);
 	return result;
 }
 
 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)