From aa681c7b99aa407d99d159924c4f2b295da3564b Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Sat, 21 Feb 2009 07:17:22 +0000 Subject: [PATCH] Fix keyword arguments for itertools.count(). Step arg without a start arg was ignored. --- Lib/test/test_itertools.py | 2 + Modules/itertoolsmodule.c | 77 +++++++++++++++++++++++--------------- 2 files changed, 48 insertions(+), 31 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 9eb138954d2..31d10dc5c84 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -350,6 +350,8 @@ class TestBasicOps(unittest.TestCase): self.assertEqual(zip('abc',count(2,3)), [('a', 2), ('b', 5), ('c', 8)]) self.assertEqual(zip('abc',count(start=2,step=3)), [('a', 2), ('b', 5), ('c', 8)]) + self.assertEqual(zip('abc',count(step=-1)), + [('a', 0), ('b', -1), ('c', -2)]) 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))) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 4aa02b62012..58b9b3219ab 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3209,19 +3209,19 @@ typedef struct { /* Counting logic and invariants: -C_add_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified. +fast_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. + When count hits Y_SSIZE_T_MAX, switch to slow_mode. -Py_add_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float. +slow_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. + Either long_cnt or long_step may be a float, Fraction, or Decimal. */ static PyTypeObject count_type; @@ -3230,6 +3230,7 @@ static PyObject * count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { countobject *lz; + int slow_mode = 0; Py_ssize_t cnt = 0; PyObject *long_cnt = NULL; PyObject *long_step = NULL; @@ -3239,36 +3240,51 @@ count_new(PyTypeObject *type, PyObject *args, PyObject *kwds) kwlist, &long_cnt, &long_step)) return NULL; - if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) || - (long_step != NULL && !PyNumber_Check(long_step))) { + 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 */ + if (cnt == -1 && PyErr_Occurred() || !PyInt_Check(long_cnt)) { PyErr_Clear(); - Py_INCREF(long_cnt); - cnt = PY_SSIZE_T_MAX; - } else - long_cnt = NULL; + slow_mode = 1; + } + Py_INCREF(long_cnt); + } else { + cnt = 0; + long_cnt = PyInt_FromLong(0); } - assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL) || - (cnt == PY_SSIZE_T_MAX && long_cnt != NULL)); + + /* If not specified, step defaults to 1 */ + if (long_step == NULL) { + long_step = PyInt_FromLong(1); + if (long_step == NULL) { + Py_DECREF(long_cnt); + return NULL; + } + } else + Py_INCREF(long_step); + + assert(long_cnt != NULL && long_step != NULL); + + /* Fast mode only works when the step is 1 */ + if (!PyInt_Check(long_step) || + PyInt_AS_LONG(long_step) != 1) { + slow_mode = 1; + } + + if (slow_mode) + cnt = PY_SSIZE_T_MAX; + else + Py_CLEAR(long_cnt); + + assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode || + cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode); + assert(slow_mode || + PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1); /* create countobject structure */ lz = (countobject *)type->tp_alloc(type, 0); @@ -3308,7 +3324,7 @@ count_nextlong(countobject *lz) long_cnt = lz->long_cnt; if (long_cnt == NULL) { - /* Switch to Py_add_mode */ + /* Switch to slow_mode */ long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX); if (long_cnt == NULL) return NULL; @@ -3360,11 +3376,10 @@ count_repr(countobject *lz) } PyDoc_STRVAR(count_doc, - "count([start[, step]]) --> count object\n\ + "count(start=0, step=1]) --> count object\n\ \n\ -Return a count object whose .next() method returns consecutive\n\ -integers starting from zero or, if specified, from start.\n\ -If step is specified, counts by that interval. Equivalent to:\n\n\ +Return a count object whose .next() method returns consecutive values.\n\ +Equivalent to:\n\n\ def count(firstval=0, step=1):\n\ x = firstval\n\ while 1:\n\