Kaydet (Commit) 73d79632 authored tarafından Raymond Hettinger's avatar Raymond Hettinger

Improve the implementation of itertools.product()

* Fix-up issues pointed-out by Neal Norwitz.
* Add extensive comments.
* The lz->result variable is now a tuple instead of a list.
* Use fast macro getitem/setitem calls so most code is in-line.
* Re-use the result tuple if available (modify in-place instead of copy).
üst c5705a82
...@@ -274,6 +274,9 @@ class TestBasicOps(unittest.TestCase): ...@@ -274,6 +274,9 @@ class TestBasicOps(unittest.TestCase):
args = map(iter, args) args = map(iter, args)
self.assertEqual(len(list(product(*args))), n) self.assertEqual(len(list(product(*args))), n)
# Test implementation detail: tuple re-use
self.assertEqual(len(set(map(id, product('abc', 'def')))), 1)
self.assertNotEqual(len(set(map(id, list(product('abc', 'def'))))), 1)
def test_repeat(self): def test_repeat(self):
self.assertEqual(zip(xrange(3),repeat('a')), self.assertEqual(zip(xrange(3),repeat('a')),
......
...@@ -1796,7 +1796,7 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds) ...@@ -1796,7 +1796,7 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
lz = (productobject *)type->tp_alloc(type, 0); lz = (productobject *)type->tp_alloc(type, 0);
if (lz == NULL) { if (lz == NULL) {
Py_DECREF(pools); Py_DECREF(pools);
return NULL; goto error;
} }
lz->pools = pools; lz->pools = pools;
...@@ -1840,7 +1840,7 @@ product_next(productobject *lz) ...@@ -1840,7 +1840,7 @@ product_next(productobject *lz)
{ {
PyObject *pool; PyObject *pool;
PyObject *elem; PyObject *elem;
PyObject *tuple_result; PyObject *oldelem;
PyObject *pools = lz->pools; PyObject *pools = lz->pools;
PyObject *result = lz->result; PyObject *result = lz->result;
Py_ssize_t npools = PyTuple_GET_SIZE(pools); Py_ssize_t npools = PyTuple_GET_SIZE(pools);
...@@ -1848,10 +1848,14 @@ product_next(productobject *lz) ...@@ -1848,10 +1848,14 @@ product_next(productobject *lz)
if (lz->stopped) if (lz->stopped)
return NULL; return NULL;
if (result == NULL) { if (result == NULL) {
/* On the first pass, return an initial tuple filled with the
first element from each pool. If any pool is empty, then
whole product is empty and we're already done */
if (npools == 0) if (npools == 0)
goto empty; goto empty;
result = PyList_New(npools); result = PyTuple_New(npools);
if (result == NULL) if (result == NULL)
goto empty; goto empty;
lz->result = result; lz->result = result;
...@@ -1861,34 +1865,61 @@ product_next(productobject *lz) ...@@ -1861,34 +1865,61 @@ product_next(productobject *lz)
goto empty; goto empty;
elem = PyTuple_GET_ITEM(pool, 0); elem = PyTuple_GET_ITEM(pool, 0);
Py_INCREF(elem); Py_INCREF(elem);
PyList_SET_ITEM(result, i, elem); PyTuple_SET_ITEM(result, i, elem);
} }
} else { } else {
Py_ssize_t *indices = lz->indices; Py_ssize_t *indices = lz->indices;
Py_ssize_t *maxvec = lz->maxvec; Py_ssize_t *maxvec = lz->maxvec;
/* Copy the previous result tuple or re-use it if available */
if (Py_REFCNT(result) > 1) {
PyObject *old_result = result;
result = PyTuple_New(npools);
if (result == NULL)
goto empty;
lz->result = result;
for (i=0; i < npools; i++) {
elem = PyTuple_GET_ITEM(old_result, i);
Py_INCREF(elem);
PyTuple_SET_ITEM(result, i, elem);
}
Py_DECREF(old_result);
}
/* Now, we've got the only copy so we can update it in-place */
assert (Py_REFCNT(result) == 1);
/* Update the pool indices right-to-left. Only advance to the
next pool when the previous one rolls-over */
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] == maxvec[i]) {
/* 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);
Py_INCREF(elem); Py_INCREF(elem);
PyList_SetItem(result, i, elem); oldelem = PyTuple_GET_ITEM(result, i);
PyTuple_SET_ITEM(result, i, elem);
Py_DECREF(oldelem);
} else { } else {
/* No rollover. Just increment and stop here. */
elem = PyTuple_GET_ITEM(pool, indices[i]); elem = PyTuple_GET_ITEM(pool, indices[i]);
Py_INCREF(elem); Py_INCREF(elem);
PyList_SetItem(result, i, elem); oldelem = PyTuple_GET_ITEM(result, i);
PyTuple_SET_ITEM(result, i, elem);
Py_DECREF(oldelem);
break; break;
} }
} }
/* If i is negative, then the indices have all rolled-over
and we're done. */
if (i < 0) if (i < 0)
return NULL; goto empty;
} }
tuple_result = PySequence_Tuple(result); Py_INCREF(result);
if (tuple_result == NULL) return result;
lz->stopped = 1;
return tuple_result;
empty: empty:
lz->stopped = 1; lz->stopped = 1;
...@@ -1898,7 +1929,7 @@ empty: ...@@ -1898,7 +1929,7 @@ empty:
PyDoc_STRVAR(product_doc, PyDoc_STRVAR(product_doc,
"product(*iterables) --> product object\n\ "product(*iterables) --> product object\n\
\n\ \n\
Cartesian product of input interables. Equivalent to nested for-loops.\n\n\ Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\ For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
The leftmost iterators are in the outermost for-loop, so the output tuples\n\ The leftmost iterators are in the outermost for-loop, so the output tuples\n\
cycle in a manner similar to an odometer (with the rightmost element changing\n\ cycle in a manner similar to an odometer (with the rightmost element changing\n\
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment