[Python-checkins] r69523 - in python/branches/py3k: Doc/library/itertools.rst Lib/test/test_itertools.py Misc/NEWS Modules/itertoolsmodule.c

Thu Feb 12 07:28:28 CET 2009

Author: raymond.hettinger
Date: Thu Feb 12 07:28:27 2009
New Revision: 69523
Log:
Issue 5032: added a step argument to itertools.count() and allowed non-integer arguments.
Modified:
 python/branches/py3k/Doc/library/itertools.rst
 python/branches/py3k/Lib/test/test_itertools.py
 python/branches/py3k/Misc/NEWS
 python/branches/py3k/Modules/itertoolsmodule.c
Modified: python/branches/py3k/Doc/library/itertools.rst
==============================================================================

--- python/branches/py3k/Doc/library/itertools.rst	(original)
+++ python/branches/py3k/Doc/library/itertools.rst	Thu Feb 12 07:28:27 2009
@@ -178,7 +178,7 @@
 
 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 @@
 # 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)
 
Modified: python/branches/py3k/Lib/test/test_itertools.py
==============================================================================
--- python/branches/py3k/Lib/test/test_itertools.py	(original)
+++ python/branches/py3k/Lib/test/test_itertools.py	Thu Feb 12 07:28:27 2009
@@ -328,7 +328,7 @@
 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 @@
 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 @@
 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('')), [])
Modified: python/branches/py3k/Misc/NEWS
==============================================================================
--- python/branches/py3k/Misc/NEWS	(original)
+++ python/branches/py3k/Misc/NEWS	Thu Feb 12 07:28:27 2009
@@ -241,6 +241,9 @@
 - 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.
Modified: python/branches/py3k/Modules/itertoolsmodule.c
==============================================================================
--- python/branches/py3k/Modules/itertoolsmodule.c	(original)
+++ python/branches/py3k/Modules/itertoolsmodule.c	Thu Feb 12 07:28:27 2009
@@ -2886,9 +2886,27 @@
 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 @@
 {
 	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 @@
 	}
 	lz->cnt = cnt;
 	lz->long_cnt = long_cnt;
+	lz->long_step = long_step;
 
 	return (PyObject *)lz;
 }
@@ -2934,7 +2970,8 @@
 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 @@
 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 @@
 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)
