[Python-checkins] r69522 - in python/trunk: Doc/library/itertools.rst Lib/test/test_itertools.py Misc/NEWS Modules/itertoolsmodule.c
raymond.hettinger
python-checkins at python.org
Thu Feb 12 06:39:47 CET 2009
Author: raymond.hettinger
Date: Thu Feb 12 06:39:46 2009
New Revision: 69522
Log:
Issue 5032: added a step argument to itertools.count() and allowed non-integer arguments.
Modified:
python/trunk/Doc/library/itertools.rst
python/trunk/Lib/test/test_itertools.py
python/trunk/Misc/NEWS
python/trunk/Modules/itertoolsmodule.c
Modified: python/trunk/Doc/library/itertools.rst
==============================================================================
--- python/trunk/Doc/library/itertools.rst (original)
+++ python/trunk/Doc/library/itertools.rst Thu Feb 12 06:39:46 2009
@@ -200,20 +200,23 @@
.. 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.
Modified: python/trunk/Lib/test/test_itertools.py
==============================================================================
--- python/trunk/Lib/test/test_itertools.py (original)
+++ python/trunk/Lib/test/test_itertools.py Thu Feb 12 06:39:46 2009
@@ -324,7 +324,7 @@
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 @@
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 @@
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('')), [])
Modified: python/trunk/Misc/NEWS
==============================================================================
--- python/trunk/Misc/NEWS (original)
+++ python/trunk/Misc/NEWS Thu Feb 12 06:39:46 2009
@@ -227,6 +227,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/trunk/Modules/itertoolsmodule.c
==============================================================================
--- python/trunk/Modules/itertoolsmodule.c (original)
+++ python/trunk/Modules/itertoolsmodule.c Thu Feb 12 06:39:46 2009
@@ -3203,9 +3203,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 *
@@ -3213,28 +3231,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 = 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 @@
}
lz->cnt = cnt;
lz->long_cnt = long_cnt;
+ lz->long_step = long_step;
return (PyObject *)lz;
}
@@ -3251,7 +3287,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);
}
@@ -3259,32 +3296,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 = 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 @@
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)
More information about the Python-checkins
mailing list