Fix keyword arguments for itertools.count().
Step arg without a start arg was ignored.
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 4aa02b6..58b9b32 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -3209,19 +3209,19 @@
/* 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 @@
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 @@
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 @@
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 @@
}
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\