diff -r 6f133133c5da Lib/test/test_itertools.py --- a/Lib/test/test_itertools.py Fri Apr 18 09:21:55 2014 -0700 +++ b/Lib/test/test_itertools.py Fri Apr 18 17:40:06 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 @@ -1087,6 +1086,24 @@ list(range(*args))) self.pickletest(islice(range(100), *args)) + # test big step + for args in [ # islice(args) should agree with range(args) + (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 + 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 6f133133c5da Modules/itertoolsmodule.c --- a/Modules/itertoolsmodule.c Fri Apr 18 09:21:55 2014 -0700 +++ b/Modules/itertoolsmodule.c Fri Apr 18 17:40:06 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,79 +1407,172 @@ 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; + PyObject *it; + PyObject *unpacked_args[3] = {NULL, NULL, 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; + Py_ssize_t i; if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds)) return NULL; - if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3)) + if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &unpacked_args[0], + &unpacked_args[1], &unpacked_args[2])) return NULL; + zero = PyLong_FromLong(0); + one = PyLong_FromLong(1); + neg_one = PyLong_FromLong(-1); + if (zero == NULL || one == NULL || neg_one == NULL) { + goto done; + } + numargs = PyTuple_Size(args); + for (i = 0; i < numargs - 1; i++) { + /* If any of the arguments cannot be stored in a Py_ssize_t, we will use + slow mode */ + if (unpacked_args[i] != Py_None && + PyLong_AsSsize_t(unpacked_args[i]) == -1 && + PyErr_Occurred()) { + PyErr_Clear(); + fast_mode = 0; + break; + } + } + if (numargs == 2) { - if (a1 != Py_None) { - stop = PyLong_AsSsize_t(a1); - if (stop == -1) { - if (PyErr_Occurred()) - PyErr_Clear(); + if (unpacked_args[0] != Py_None) { + int is_bad; + if (fast_mode) { + values.stop = PyLong_AsSsize_t(unpacked_args[0]); + is_bad = values.stop < 0; + } else { + Py_INCREF(zero); + long_values.next = zero; + Py_INCREF(one); + long_values.step = one; + Py_INCREF(unpacked_args[0]); + long_values.stop = unpacked_args[0]; + is_bad = PyObject_RichCompareBool(long_values.stop, zero, Py_GE) <= 0; + } + if (is_bad) { 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; } + } else if (!fast_mode) { + Py_INCREF(zero); + long_values.next = zero; + Py_INCREF(Py_None); + long_values.stop = Py_None; + Py_INCREF(one); + long_values.step = one; } } else { - if (a1 != Py_None) - start = PyLong_AsSsize_t(a1); - if (start == -1 && PyErr_Occurred()) - PyErr_Clear(); - if (a2 != Py_None) { - stop = PyLong_AsSsize_t(a2); - if (stop == -1) { - if (PyErr_Occurred()) - PyErr_Clear(); + int is_bad; + + if (unpacked_args[0] != Py_None) { + if (fast_mode) { + values.next = PyLong_AsSsize_t(unpacked_args[0]); + is_bad = values.next < 0; + } else { + Py_INCREF(unpacked_args[0]); + long_values.next = unpacked_args[0]; + is_bad = PyObject_RichCompareBool(long_values.next, zero, Py_GE) < 0; + } + if (is_bad) { PyErr_SetString(PyExc_ValueError, - "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize."); - return NULL; + "Start argument for islice() must be None or a non-negative integer."); + goto done; } + } else if (fast_mode) { + values.next = 0; + } else { + Py_INCREF(zero); + long_values.next = zero; + } + + if (unpacked_args[1] != Py_None) { + if (fast_mode) { + values.stop = PyLong_AsSsize_t(unpacked_args[1]); + is_bad = values.stop < 0; + } else { + Py_INCREF(unpacked_args[1]); + long_values.stop = unpacked_args[1]; + is_bad = PyObject_RichCompareBool(long_values.stop, zero, Py_GE) <= 0; + } + if (is_bad) { + PyErr_SetString(PyExc_ValueError, + "Stop argument for islice() must be None or a non-negative integer."); + goto done; + } + } else { + Py_INCREF(Py_None); + long_values.stop = Py_None; + } + + if (unpacked_args[2] && unpacked_args[2] != Py_None) { + if (fast_mode) { + values.step = PyLong_AsSsize_t(unpacked_args[2]); + is_bad = values.step <= 0; + } else { + Py_INCREF(unpacked_args[2]); + long_values.step = unpacked_args[2]; + is_bad = PyObject_RichCompareBool(long_values.step, zero, Py_GT) <= 0; + } + if (is_bad) { + PyErr_SetString(PyExc_ValueError, + "Step argument for islice() must be None or a positive number."); + goto done; + } + } else { + Py_INCREF(one); + long_values.step = one; } } - 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; + + if (fast_mode) { + values.cnt = 0; + } else { + 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 +1582,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); } @@ -1484,34 +1597,117 @@ } 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); + return NULL; + } + oldnext = next; + next = lz->long_values.next = PyNumber_Add(next, step); + if (lz->long_values.next == NULL) { + Py_DECREF(one); + return NULL; + } + 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(one); + return item; +} + +static PyObject * islice_next(isliceobject *lz) { PyObject *item; PyObject *it = lz->it; - Py_ssize_t stop = lz->stop; + int fast_mode = lz->fast_mode; + 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 (fast_mode && stop == -1 && (lz->values.next > PY_SSIZE_T_MAX - lz->values.step)) { + fast_mode = lz->fast_mode = 0; + + lz->long_values.cnt = PyLong_FromSsize_t(lz->values.cnt); + lz->long_values.step = PyLong_FromSsize_t(lz->values.step); + lz->long_values.next = PyLong_FromSsize_t(lz->values.next); + Py_INCREF(Py_None); + lz->long_values.stop = Py_None; + } + + if (!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; } @@ -1522,26 +1718,39 @@ * then 'setstate' with the next and count */ PyObject *stop; - if (lz->stop == -1) { - stop = Py_None; - Py_INCREF(stop); + if (lz->fast_mode) { + if (lz->values.stop == -1) { + stop = Py_None; + Py_INCREF(stop); + } else { + stop = PyLong_FromSsize_t(lz->values.stop); + if (stop == NULL) + return NULL; + } + return Py_BuildValue("O(OnNn)n", Py_TYPE(lz), + lz->it, lz->values.next, stop, lz->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; + if (lz->fast_mode) { + cnt = PyLong_AsSsize_t(state); + if (cnt == -1 && PyErr_Occurred()) + return NULL; + + lz->values.cnt = cnt; + } else { + Py_INCREF(state); + lz->long_values.cnt = state; + } Py_RETURN_NONE; }