mirrors/cpython
1
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2025-11-27 13:45:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 10

Handle the repeat keyword argument for itertools.product().

Browse source
This commit is contained in:
Raymond Hettinger 2008-02-29 02:21:48 +00:00
parent 1ddf1d8482
commit 08ff6822cc
2 changed files with 30 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

3
Lib/test/test_itertools.py
View file

@ -296,6 +296,9 @@ class TestBasicOps(unittest.TestCase):
([range(2), range(3), range(0)], []), # last iterable with zero length ([range(2), range(3), range(0)], []), # last iterable with zero length
]: ]:
self.assertEqual(list(product(*args)), result) self.assertEqual(list(product(*args)), result)
for r in range(4):
self.assertEqual(list(product(*(args*r))),
list(product(*args, **dict(repeat=r))))
self.assertEqual(len(list(product(*[range(7)]*6))), 7**6) self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
self.assertRaises(TypeError, product, range(6), None) self.assertRaises(TypeError, product, range(6), None)
argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3), argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),

32
Modules/itertoolsmodule.c
View file

@ -1782,17 +1782,32 @@ static PyObject *
product_new(PyTypeObject *type, PyObject *args, PyObject *kwds) product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{ {
productobject *lz; productobject *lz;
Py_ssize_t npools; Py_ssize_t nargs, npools, repeat=1;
PyObject *pools = NULL; PyObject *pools = NULL;
Py_ssize_t *maxvec = NULL; Py_ssize_t *maxvec = NULL;
Py_ssize_t *indices = NULL; Py_ssize_t *indices = NULL;
Py_ssize_t i; Py_ssize_t i;
if (type == &product_type && !_PyArg_NoKeywords("product()", kwds)) if (kwds != NULL) {
return NULL; char *kwlist[] = {"repeat", 0};
PyObject *tmpargs = PyTuple_New(0);
if (tmpargs == NULL)
return NULL;
if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
Py_DECREF(tmpargs);
return NULL;
}
Py_DECREF(tmpargs);
if (repeat < 0) {
PyErr_SetString(PyExc_ValueError,
"repeat argument cannot be negative");
return NULL;
}
}
assert(PyTuple_Check(args)); assert(PyTuple_Check(args));
npools = PyTuple_GET_SIZE(args); nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
npools = nargs * repeat;
maxvec = PyMem_Malloc(npools * sizeof(Py_ssize_t)); maxvec = PyMem_Malloc(npools * sizeof(Py_ssize_t));
indices = PyMem_Malloc(npools * sizeof(Py_ssize_t)); indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
@ -1805,7 +1820,7 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
if (pools == NULL) if (pools == NULL)
goto error; goto error;
for (i=0; i < npools; ++i) { for (i=0; i < nargs ; ++i) {
PyObject *item = PyTuple_GET_ITEM(args, i); PyObject *item = PyTuple_GET_ITEM(args, i);
PyObject *pool = PySequence_Tuple(item); PyObject *pool = PySequence_Tuple(item);
if (pool == NULL) if (pool == NULL)
@ -1815,6 +1830,13 @@ product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
maxvec[i] = PyTuple_GET_SIZE(pool); maxvec[i] = PyTuple_GET_SIZE(pool);
indices[i] = 0; indices[i] = 0;
} }
for ( ; i < npools; ++i) {
PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
Py_INCREF(pool);
PyTuple_SET_ITEM(pools, i, pool);
maxvec[i] = maxvec[i - nargs];
indices[i] = 0;
}
/* create productobject structure */ /* create productobject structure */
lz = (productobject *)type->tp_alloc(type, 0); lz = (productobject *)type->tp_alloc(type, 0);