diff -r 8c8315bac6a8 Lib/test/test_itertools.py --- a/Lib/test/test_itertools.py Sun Apr 20 09:45:00 2014 -0700 +++ b/Lib/test/test_itertools.py Mon Apr 21 16:04:36 2014 -0700 @@ -10,7 +10,6 @@ import copy import pickle from functools import reduce -import sys import struct maxsize = support.MAX_Py_ssize_t minsize = -maxsize-1 @@ -1079,7 +1078,14 @@ (10, 3, 20), (10, 20), (10, 3), - (20,) + (20,), + (0, 10, sys.maxsize+1), + (10, 20, sys.maxsize+1), + (10, 3, sys.maxsize+1), + (10, 0, sys.maxsize+1), + (sys.maxsize+1, 0, 1), + (sys.maxsize+1, sys.maxsize, 1), + (sys.maxsize+1, 0), ]: self.assertEqual(list(copy.copy(islice(range(100), *args))), list(range(*args))) @@ -1087,6 +1093,25 @@ list(range(*args))) self.pickletest(islice(range(100), *args)) + # test big step + for args in [ # islice(args) should agree with range(args) + (0, 10, sys.maxsize+1), + (10, 20, sys.maxsize+1), + (10, 3, sys.maxsize+1), + (10, 0, sys.maxsize+1), + ]: + self.assertEqual(list(islice(range(100), *args)), + list(range(*args))) + + # test big start / stop + for args in [ + (sys.maxsize+1, 0, 1), + (sys.maxsize+1, sys.maxsize, 1), + (sys.maxsize+1, 0), + ]: + self.assertEqual(list(islice(range(100), *args)), + list(range(*args))) + def test_takewhile(self): data = [1, 3, 5, 20, 2, 4, 6, 8] self.assertEqual(list(takewhile(underten, data)), [1, 3, 5]) diff -r 8c8315bac6a8 Modules/itertoolsmodule.c --- a/Modules/itertoolsmodule.c Sun Apr 20 09:45:00 2014 -0700 +++ b/Modules/itertoolsmodule.c Mon Apr 21 16:04:36 2014 -0700 @@ -1376,13 +1376,29 @@ /* islice object ************************************************************/ -typedef struct { - PyObject_HEAD - PyObject *it; +/* islice_values is used for fast iteration when all of 'start', 'stop', and + 'step' arguments fit in Py_ssize_t */ +struct islice_values { Py_ssize_t next; Py_ssize_t stop; Py_ssize_t step; Py_ssize_t cnt; +}; + +/* islice_long_values is used for slower iteration when the numbers are big */ +struct islice_long_values { + PyObject *next; + PyObject *stop; + PyObject *step; + PyObject *cnt; +}; + +typedef struct { + PyObject_HEAD + PyObject *it; + int fast_mode; + struct islice_long_values long_values; + struct islice_values values; } isliceobject; static PyTypeObject islice_type; @@ -1391,10 +1407,15 @@ islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { PyObject *seq; - Py_ssize_t start=0, stop=-1, step=1; PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL; Py_ssize_t numargs; - isliceobject *lz; + isliceobject *lz = NULL; + struct islice_long_values long_values = {NULL, NULL, NULL, NULL}; + struct islice_values values = {0, -1, 1, 0}; + PyObject *zero = NULL; + PyObject *one = NULL; + PyObject *neg_one = NULL; + int fast_mode = 1; if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds)) return NULL; @@ -1403,67 +1424,126 @@ return NULL; numargs = PyTuple_Size(args); + zero = PyLong_FromLong(0); + one = PyLong_FromLong(1); + neg_one = PyLong_FromLong(-1); + if (zero == NULL || one == NULL || neg_one == NULL) { + goto done; + } + if (numargs == 2) { if (a1 != Py_None) { - stop = PyLong_AsSsize_t(a1); - if (stop == -1) { - if (PyErr_Occurred()) - PyErr_Clear(); + if (PyObject_RichCompareBool(a1, zero, Py_GE) <= 0) { PyErr_SetString(PyExc_ValueError, - "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize."); - return NULL; + "Stop argument for islice() must be None or a non-negative integer."); + goto done; + } + values.stop = PyLong_AsSsize_t(a1); + if (values.stop == -1 && PyErr_Occurred()) { + /* 'stop' cannot be represented in Py_ssize_t. We might still + be able to continue in fast mode until we encounter an + overflow. */ + PyErr_Clear(); } } + Py_INCREF(zero); + long_values.next = zero; + Py_INCREF(one); + long_values.step = one; + Py_INCREF(a1); + long_values.stop = a1; } else { - if (a1 != Py_None) - start = PyLong_AsSsize_t(a1); - if (start == -1 && PyErr_Occurred()) - PyErr_Clear(); + if (a1 != Py_None) { + if (PyObject_RichCompareBool(a1, zero, Py_GE) <= 0) { + PyErr_SetString(PyExc_ValueError, + "Start argument for islice() must be None or a non-negative integer."); + goto done; + } + Py_INCREF(a1); + long_values.next = a1; + values.next = PyLong_AsSsize_t(a1); + if (values.next == -1 && PyErr_Occurred()) { + /* slicing starts at an index that cannot be represented by + Py_ssize_t. We will use "slow mode" for iteration. */ + fast_mode = 0; + PyErr_Clear(); + } + } else { + Py_INCREF(zero); + long_values.next = zero; + values.next = 0; + } + if (a2 != Py_None) { - stop = PyLong_AsSsize_t(a2); - if (stop == -1) { - if (PyErr_Occurred()) - PyErr_Clear(); + if (PyObject_RichCompareBool(a2, zero, Py_GE) <= 0) { PyErr_SetString(PyExc_ValueError, - "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize."); - return NULL; + "Stop argument for islice() must be None or a non-negative integer."); + goto done; } + Py_INCREF(a2); + long_values.stop = a2; + values.stop = PyLong_AsSsize_t(a2); + if (values.stop == -1 && PyErr_Occurred()) { + PyErr_Clear(); + } + } else { + Py_INCREF(Py_None); + long_values.stop = Py_None; + values.stop = -1; + } + + if (a3 && a3 != Py_None) { + if (PyObject_RichCompareBool(a3, zero, Py_GT) <= 0) { + PyErr_SetString(PyExc_ValueError, + "Step argument for islice() must be None or a positive number."); + goto done; + } + Py_INCREF(a3); + long_values.step = a3; + values.step = PyLong_AsSsize_t(a3); + if (values.step == -1 && PyErr_Occurred()) { + PyErr_Clear(); + } + } else { + Py_INCREF(one); + long_values.step = one; + values.step = 1; } } - if (start<0 || stop<-1) { - PyErr_SetString(PyExc_ValueError, - "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize."); - return NULL; - } - - if (a3 != NULL) { - if (a3 != Py_None) - step = PyLong_AsSsize_t(a3); - if (step == -1 && PyErr_Occurred()) - PyErr_Clear(); - } - if (step<1) { - PyErr_SetString(PyExc_ValueError, - "Step for islice() must be a positive integer or None."); - return NULL; - } + + values.cnt = 0; + Py_INCREF(zero); + long_values.cnt = zero; /* Get iterator. */ it = PyObject_GetIter(seq); if (it == NULL) - return NULL; + goto done; /* create isliceobject structure */ lz = (isliceobject *)type->tp_alloc(type, 0); if (lz == NULL) { Py_DECREF(it); + goto done; + } + lz->it = it; + lz->fast_mode = fast_mode; + lz->long_values = long_values; + lz->values = values; + +done: + Py_DECREF(zero); + Py_DECREF(one); + Py_DECREF(neg_one); + + if (!lz) { + Py_XDECREF(long_values.next); + Py_XDECREF(long_values.stop); + Py_XDECREF(long_values.step); + Py_XDECREF(long_values.cnt); + return NULL; } - lz->it = it; - lz->next = start; - lz->stop = stop; - lz->step = step; - lz->cnt = 0L; return (PyObject *)lz; } @@ -1473,6 +1553,10 @@ { PyObject_GC_UnTrack(lz); Py_XDECREF(lz->it); + Py_XDECREF(lz->long_values.next); + Py_XDECREF(lz->long_values.stop); + Py_XDECREF(lz->long_values.step); + Py_XDECREF(lz->long_values.cnt); Py_TYPE(lz)->tp_free(lz); } @@ -1480,38 +1564,132 @@ islice_traverse(isliceobject *lz, visitproc visit, void *arg) { Py_VISIT(lz->it); + Py_VISIT(lz->long_values.next); + Py_VISIT(lz->long_values.stop); + Py_VISIT(lz->long_values.step); + Py_VISIT(lz->long_values.cnt); return 0; } static PyObject * +islice_next_slow(isliceobject *lz) +{ + PyObject *item; + PyObject *it = lz->it; + PyObject *stop, *next, *step; + PyObject *oldnext; + PyObject *one; + PyObject *tmp; + PyObject *(*iternext)(PyObject *); + + one = PyLong_FromLong(1); + if (!one) { + return NULL; + } + stop = lz->long_values.stop; + next = lz->long_values.next; + step = lz->long_values.step; + + iternext = *Py_TYPE(it)->tp_iternext; + while (PyObject_RichCompareBool(lz->long_values.cnt, next, Py_LT) == 1) { + item = iternext(it); + if (item == NULL) { + Py_DECREF(one); + return NULL; + } + Py_DECREF(item); + tmp = PyNumber_Add(lz->long_values.cnt, one); + if (tmp == NULL) { + Py_DECREF(one); + return NULL; + } + Py_DECREF(lz->long_values.cnt); + lz->long_values.cnt = tmp; + } + + if (stop != Py_None && + PyObject_RichCompareBool(lz->long_values.cnt, stop, Py_GE) == 1) { + Py_DECREF(one); + return NULL; + } + item = iternext(it); + if (item == NULL) { + Py_DECREF(one); + return NULL; + } + + tmp = PyNumber_Add(lz->long_values.cnt, one); + if (tmp == NULL) { + Py_DECREF(one); + Py_DECREF(item); + return NULL; + } + Py_DECREF(lz->long_values.cnt); + lz->long_values.cnt = tmp; + + oldnext = next; + next = PyNumber_Add(next, step); + if (next == NULL) { + Py_DECREF(one); + Py_DECREF(item); + return NULL; + } + lz->long_values.next = next; + + if (PyObject_RichCompareBool(next, oldnext, Py_LT) == 1 || + (stop != Py_None && PyObject_RichCompareBool(next, stop, Py_GT) == 1)) { + Py_DECREF(lz->long_values.next); + Py_INCREF(stop); + lz->long_values.next = stop; + } + Py_DECREF(oldnext); + Py_DECREF(one); + return item; +} + +static PyObject * islice_next(isliceobject *lz) { PyObject *item; PyObject *it = lz->it; - Py_ssize_t stop = lz->stop; + Py_ssize_t stop = lz->values.stop; Py_ssize_t oldnext; PyObject *(*iternext)(PyObject *); + /* Switch to slow mode if there would be an overflow */ + if (lz->fast_mode && (lz->values.next > PY_SSIZE_T_MAX - lz->values.step)) { + lz->fast_mode = 0; + + Py_XDECREF(lz->long_values.cnt); + lz->long_values.cnt = PyLong_FromSsize_t(lz->values.cnt); + Py_XDECREF(lz->long_values.next); + lz->long_values.next = PyLong_FromSsize_t(lz->values.next); + } + + if (!lz->fast_mode) { + return islice_next_slow(lz); + } + iternext = *Py_TYPE(it)->tp_iternext; - while (lz->cnt < lz->next) { + while (lz->values.cnt < lz->values.next) { item = iternext(it); if (item == NULL) return NULL; Py_DECREF(item); - lz->cnt++; + lz->values.cnt++; } - if (stop != -1 && lz->cnt >= stop) + if (stop != -1 && lz->values.cnt >= stop) return NULL; item = iternext(it); if (item == NULL) return NULL; - lz->cnt++; - oldnext = lz->next; + lz->values.cnt++; + oldnext = lz->values.next; /* The (size_t) cast below avoids the danger of undefined behaviour from signed integer overflow. */ - lz->next += (size_t)lz->step; - if (lz->next < oldnext || (stop != -1 && lz->next > stop)) - lz->next = stop; + lz->values.next += (size_t)lz->values.step; + if (lz->values.next < oldnext || (stop != -1 && lz->values.next > stop)) + lz->values.next = stop; return item; } @@ -1521,27 +1699,33 @@ /* When unpickled, generate a new object with the same bounds, * then 'setstate' with the next and count */ - PyObject *stop; - if (lz->stop == -1) { - stop = Py_None; - Py_INCREF(stop); + if (lz->fast_mode) { + /* We could be in "fast mode" but with 'step' or 'stop' that are outside + the range of a Py_ssize_t, so we use long_values.step and + long_values.step. */ + return Py_BuildValue("O(OnOO)n", Py_TYPE(lz), + lz->it, lz->values.next, lz->long_values.stop, + lz->long_values.step, lz->values.cnt); } else { - stop = PyLong_FromSsize_t(lz->stop); - if (stop == NULL) - return NULL; + return Py_BuildValue("O(OOOO)O", Py_TYPE(lz), + lz->it, lz->long_values.next, lz->long_values.stop, + lz->long_values.step, lz->long_values.cnt); } - return Py_BuildValue("O(OnNn)n", Py_TYPE(lz), - lz->it, lz->next, stop, lz->step, - lz->cnt); } static PyObject * islice_setstate(isliceobject *lz, PyObject *state) { - Py_ssize_t cnt = PyLong_AsSsize_t(state); - if (cnt == -1 && PyErr_Occurred()) - return NULL; - lz->cnt = cnt; + Py_ssize_t cnt; + cnt = PyLong_AsSsize_t(state); + if (cnt == -1 && PyErr_Occurred()) { + PyErr_Clear(); + lz->fast_mode = 0; + } + + lz->values.cnt = cnt; + Py_INCREF(state); + lz->long_values.cnt = state; Py_RETURN_NONE; }