Simplify code for itertools.product().

This commit is contained in:
Raymond Hettinger 2008-03-02 11:57:16 +00:00
parent 2b7a5c468b
commit 61024b9181
1 changed files with 2 additions and 13 deletions

View File

@ -1770,7 +1770,6 @@ static PyTypeObject chain_type = {
typedef struct { typedef struct {
PyObject_HEAD PyObject_HEAD
PyObject *pools; /* tuple of pool tuples */ PyObject *pools; /* tuple of pool tuples */
Py_ssize_t *maxvec; /* size of each pool */
Py_ssize_t *indices; /* one index per pool */ Py_ssize_t *indices; /* one index per pool */
PyObject *result; /* most recently returned result tuple */ PyObject *result; /* most recently returned result tuple */
int stopped; /* set to 1 when the product iterator is exhausted */ int stopped; /* set to 1 when the product iterator is exhausted */
@ -1784,7 +1783,6 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
productobject *lz; productobject *lz;
Py_ssize_t nargs, npools, repeat=1; Py_ssize_t nargs, npools, repeat=1;
PyObject *pools = NULL; PyObject *pools = NULL;
Py_ssize_t *maxvec = NULL;
Py_ssize_t *indices = NULL; Py_ssize_t *indices = NULL;
Py_ssize_t i; Py_ssize_t i;
@ -1809,9 +1807,8 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args); nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
npools = nargs * repeat; npools = nargs * repeat;
maxvec = PyMem_Malloc(npools * sizeof(Py_ssize_t));
indices = PyMem_Malloc(npools * sizeof(Py_ssize_t)); indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
if (maxvec == NULL || indices == NULL) { if (indices == NULL) {
PyErr_NoMemory(); PyErr_NoMemory();
goto error; goto error;
} }
@ -1825,16 +1822,13 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
PyObject *pool = PySequence_Tuple(item); PyObject *pool = PySequence_Tuple(item);
if (pool == NULL) if (pool == NULL)
goto error; goto error;
PyTuple_SET_ITEM(pools, i, pool); PyTuple_SET_ITEM(pools, i, pool);
maxvec[i] = PyTuple_GET_SIZE(pool);
indices[i] = 0; indices[i] = 0;
} }
for ( ; i < npools; ++i) { for ( ; i < npools; ++i) {
PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs); PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
Py_INCREF(pool); Py_INCREF(pool);
PyTuple_SET_ITEM(pools, i, pool); PyTuple_SET_ITEM(pools, i, pool);
maxvec[i] = maxvec[i - nargs];
indices[i] = 0; indices[i] = 0;
} }
@ -1844,7 +1838,6 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
goto error; goto error;
lz->pools = pools; lz->pools = pools;
lz->maxvec = maxvec;
lz->indices = indices; lz->indices = indices;
lz->result = NULL; lz->result = NULL;
lz->stopped = 0; lz->stopped = 0;
@ -1852,8 +1845,6 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
return (PyObject *)lz; return (PyObject *)lz;
error: error:
if (maxvec != NULL)
PyMem_Free(maxvec);
if (indices != NULL) if (indices != NULL)
PyMem_Free(indices); PyMem_Free(indices);
Py_XDECREF(pools); Py_XDECREF(pools);
@ -1866,7 +1857,6 @@ product_dealloc(productobject *lz)
PyObject_GC_UnTrack(lz); PyObject_GC_UnTrack(lz);
Py_XDECREF(lz->pools); Py_XDECREF(lz->pools);
Py_XDECREF(lz->result); Py_XDECREF(lz->result);
PyMem_Free(lz->maxvec);
PyMem_Free(lz->indices); PyMem_Free(lz->indices);
Py_TYPE(lz)->tp_free(lz); Py_TYPE(lz)->tp_free(lz);
} }
@ -1913,7 +1903,6 @@ product_next(productobject *lz)
} }
} else { } else {
Py_ssize_t *indices = lz->indices; Py_ssize_t *indices = lz->indices;
Py_ssize_t *maxvec = lz->maxvec;
/* Copy the previous result tuple or re-use it if available */ /* Copy the previous result tuple or re-use it if available */
if (Py_REFCNT(result) > 1) { if (Py_REFCNT(result) > 1) {
@ -1937,7 +1926,7 @@ product_next(productobject *lz)
for (i=npools-1 ; i >= 0 ; i--) { for (i=npools-1 ; i >= 0 ; i--) {
pool = PyTuple_GET_ITEM(pools, i); pool = PyTuple_GET_ITEM(pools, i);
indices[i]++; indices[i]++;
if (indices[i] == maxvec[i]) { if (indices[i] == PyTuple_GET_SIZE(pool)) {
/* Roll-over and advance to next pool */ /* Roll-over and advance to next pool */
indices[i] = 0; indices[i] = 0;
elem = PyTuple_GET_ITEM(pool, 0); elem = PyTuple_GET_ITEM(pool, 0);