diff options
Diffstat (limited to 'Modules/itertoolsmodule.c')
-rw-r--r-- | Modules/itertoolsmodule.c | 121 |
1 files changed, 98 insertions, 23 deletions
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 57a398f..4693fbc 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -401,7 +401,7 @@ static PyTypeObject _grouper_type = { typedef struct { PyObject_HEAD PyObject *it; - int numread; + int numread; /* 0 <= numread <= LINKCELLS */ PyObject *nextlink; PyObject *(values[LINKCELLS]); } teedataobject; @@ -409,7 +409,7 @@ typedef struct { typedef struct { PyObject_HEAD teedataobject *dataobj; - int index; + int index; /* 0 <= index <= LINKCELLS */ PyObject *weakreflist; } teeobject; @@ -1492,19 +1492,22 @@ islice_next(isliceobject *lz) Py_ssize_t oldnext; PyObject *(*iternext)(PyObject *); + if (it == NULL) + return NULL; + iternext = *Py_TYPE(it)->tp_iternext; while (lz->cnt < lz->next) { item = iternext(it); if (item == NULL) - return NULL; + goto empty; Py_DECREF(item); lz->cnt++; } if (stop != -1 && lz->cnt >= stop) - return NULL; + goto empty; item = iternext(it); if (item == NULL) - return NULL; + goto empty; lz->cnt++; oldnext = lz->next; /* The (size_t) cast below avoids the danger of undefined @@ -1513,6 +1516,10 @@ islice_next(isliceobject *lz) if (lz->next < oldnext || (stop != -1 && lz->next > stop)) lz->next = stop; return item; + +empty: + Py_CLEAR(lz->it); + return NULL; } static PyObject * @@ -1522,6 +1529,18 @@ islice_reduce(isliceobject *lz) * then 'setstate' with the next and count */ PyObject *stop; + if (lz->it == NULL) { + PyObject *empty_list; + PyObject *empty_it; + empty_list = PyList_New(0); + if (empty_list == NULL) + return NULL; + empty_it = PyObject_GetIter(empty_list); + Py_DECREF(empty_list); + if (empty_it == NULL) + return NULL; + return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0); + } if (lz->stop == -1) { stop = Py_None; Py_INCREF(stop); @@ -2065,6 +2084,18 @@ product_dealloc(productobject *lz) Py_TYPE(lz)->tp_free(lz); } +static PyObject * +product_sizeof(productobject *lz, void *unused) +{ + Py_ssize_t res; + + res = sizeof(productobject); + res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t); + return PyLong_FromSsize_t(res); +} + +PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes."); + static int product_traverse(productobject *lz, visitproc visit, void *arg) { @@ -2242,6 +2273,8 @@ static PyMethodDef product_methods[] = { reduce_doc}, {"__setstate__", (PyCFunction)product_setstate, METH_O, setstate_doc}, + {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -2382,6 +2415,16 @@ combinations_dealloc(combinationsobject *co) Py_TYPE(co)->tp_free(co); } +static PyObject * +combinations_sizeof(combinationsobject *co, void *unused) +{ + Py_ssize_t res; + + res = sizeof(combinationsobject); + res += co->r * sizeof(Py_ssize_t); + return PyLong_FromSsize_t(res); +} + static int combinations_traverse(combinationsobject *co, visitproc visit, void *arg) { @@ -2553,6 +2596,8 @@ static PyMethodDef combinations_methods[] = { reduce_doc}, {"__setstate__", (PyCFunction)combinations_setstate, METH_O, setstate_doc}, + {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -2711,6 +2756,16 @@ cwr_dealloc(cwrobject *co) Py_TYPE(co)->tp_free(co); } +static PyObject * +cwr_sizeof(cwrobject *co, void *unused) +{ + Py_ssize_t res; + + res = sizeof(cwrobject); + res += co->r * sizeof(Py_ssize_t); + return PyLong_FromSsize_t(res); +} + static int cwr_traverse(cwrobject *co, visitproc visit, void *arg) { @@ -2740,11 +2795,13 @@ cwr_next(cwrobject *co) if (result == NULL) goto empty; co->result = result; - elem = PyTuple_GET_ITEM(pool, 0); - for (i=0; i<r ; i++) { - assert(indices[i] == 0); - Py_INCREF(elem); - PyTuple_SET_ITEM(result, i, elem); + if (n > 0) { + elem = PyTuple_GET_ITEM(pool, 0); + for (i=0; i<r ; i++) { + assert(indices[i] == 0); + Py_INCREF(elem); + PyTuple_SET_ITEM(result, i, elem); + } } } else { /* Copy the previous result tuple or re-use it if available */ @@ -2870,6 +2927,8 @@ static PyMethodDef cwr_methods[] = { reduce_doc}, {"__setstate__", (PyCFunction)cwr_setstate, METH_O, setstate_doc}, + {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -2936,18 +2995,18 @@ def permutations(iterable, r=None): cycles = range(n-r+1, n+1)[::-1] yield tuple(pool[i] for i in indices[:r]) while n: - for i in reversed(range(r)): - cycles[i] -= 1 - if cycles[i] == 0: - indices[i:] = indices[i+1:] + indices[i:i+1] - cycles[i] = n - i + for i in reversed(range(r)): + cycles[i] -= 1 + if cycles[i] == 0: + indices[i:] = indices[i+1:] + indices[i:i+1] + cycles[i] = n - i + else: + j = cycles[i] + indices[i], indices[-j] = indices[-j], indices[i] + yield tuple(pool[i] for i in indices[:r]) + break else: - j = cycles[i] - indices[i], indices[-j] = indices[-j], indices[i] - yield tuple(pool[i] for i in indices[:r]) - break - else: - return + return */ typedef struct { @@ -3046,6 +3105,17 @@ permutations_dealloc(permutationsobject *po) Py_TYPE(po)->tp_free(po); } +static PyObject * +permutations_sizeof(permutationsobject *po, void *unused) +{ + Py_ssize_t res; + + res = sizeof(permutationsobject); + res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t); + res += po->r * sizeof(Py_ssize_t); + return PyLong_FromSsize_t(res); +} + static int permutations_traverse(permutationsobject *po, visitproc visit, void *arg) { @@ -3251,6 +3321,8 @@ static PyMethodDef permuations_methods[] = { reduce_doc}, {"__setstate__", (PyCFunction)permutations_setstate, METH_O, setstate_doc}, + {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -4055,14 +4127,17 @@ repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { repeatobject *ro; PyObject *element; - Py_ssize_t cnt = -1; + Py_ssize_t cnt = -1, n_kwds = 0; static char *kwargs[] = {"object", "times", NULL}; if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs, &element, &cnt)) return NULL; - if (PyTuple_Size(args) == 2 && cnt < 0) + if (kwds != NULL) + n_kwds = PyDict_Size(kwds); + /* Does user supply times argument? */ + if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0) cnt = 0; ro = (repeatobject *)type->tp_alloc(type, 0); |