mirror of https://github.com/python/cpython.git
First draft for itertools.product(). Docs and other updates forthcoming.
This commit is contained in:
Raymond Hettinger
parent
12db865a64
commit
50986cc45b
|
|
@ -253,6 +253,28 @@ def test_iziplongest(self):
|
||||||
ids = map(id, list(izip_longest('abc', 'def')))
|
ids = map(id, list(izip_longest('abc', 'def')))
|
||||||
self.assertEqual(len(dict.fromkeys(ids)), len(ids))
|
self.assertEqual(len(dict.fromkeys(ids)), len(ids))
|
||||||
|
|
||||||
|
def test_product(self):
|
||||||
|
for args, result in [
|
||||||
|
([], []), # zero iterables ??? is this correct
|
||||||
|
(['ab'], [('a',), ('b',)]), # one iterable
|
||||||
|
([range(2), range(3)], [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]), # two iterables
|
||||||
|
([range(0), range(2), range(3)], []), # first iterable with zero length
|
||||||
|
([range(2), range(0), range(3)], []), # middle iterable with zero length
|
||||||
|
([range(2), range(3), range(0)], []), # last iterable with zero length
|
||||||
|
]:
|
||||||
|
self.assertEqual(list(product(*args)), result)
|
||||||
|
self.assertEqual(len(list(product(*[range(7)]*6))), 7**6)
|
||||||
|
self.assertRaises(TypeError, product, range(6), None)
|
||||||
|
argtypes = ['', 'abc', '', xrange(0), xrange(4), dict(a=1, b=2, c=3),
|
||||||
|
set('abcdefg'), range(11), tuple(range(13))]
|
||||||
|
for i in range(100):
|
||||||
|
args = [random.choice(argtypes) for j in range(random.randrange(5))]
|
||||||
|
n = reduce(operator.mul, map(len, args), 1) if args else 0
|
||||||
|
self.assertEqual(len(list(product(*args))), n)
|
||||||
|
args = map(iter, args)
|
||||||
|
self.assertEqual(len(list(product(*args))), n)
|
||||||
|
|
||||||
|
|
||||||
def test_repeat(self):
|
def test_repeat(self):
|
||||||
self.assertEqual(zip(xrange(3),repeat('a')),
|
self.assertEqual(zip(xrange(3),repeat('a')),
|
||||||
[(0, 'a'), (1, 'a'), (2, 'a')])
|
[(0, 'a'), (1, 'a'), (2, 'a')])
|
||||||
|
|
@ -623,6 +645,12 @@ def test_chain(self):
|
||||||
self.assertRaises(TypeError, list, chain(N(s)))
|
self.assertRaises(TypeError, list, chain(N(s)))
|
||||||
self.assertRaises(ZeroDivisionError, list, chain(E(s)))
|
self.assertRaises(ZeroDivisionError, list, chain(E(s)))
|
||||||
|
|
||||||
|
def test_product(self):
|
||||||
|
for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
|
||||||
|
self.assertRaises(TypeError, product, X(s))
|
||||||
|
self.assertRaises(TypeError, product, N(s))
|
||||||
|
self.assertRaises(ZeroDivisionError, product, E(s))
|
||||||
|
|
||||||
def test_cycle(self):
|
def test_cycle(self):
|
||||||
for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
|
for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
|
||||||
for g in (G, I, Ig, S, L, R):
|
for g in (G, I, Ig, S, L, R):
|
||||||
|
|
|
||||||
|
|
@ -1741,6 +1741,216 @@ static PyTypeObject chain_type = {
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
||||||
|
/* product object ************************************************************/
|
||||||
|
|
||||||
|
typedef struct {
|
||||||
|
PyObject_HEAD
|
||||||
|
PyObject *pools; /* tuple of pool tuples */
|
||||||
|
Py_ssize_t *maxvec;
|
||||||
|
Py_ssize_t *indices;
|
||||||
|
PyObject *result;
|
||||||
|
int stopped;
|
||||||
|
} productobject;
|
||||||
|
|
||||||
|
static PyTypeObject product_type;
|
||||||
|
|
||||||
|
static PyObject *
|
||||||
|
product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
||||||
|
{
|
||||||
|
productobject *lz;
|
||||||
|
Py_ssize_t npools;
|
||||||
|
PyObject *pools = NULL;
|
||||||
|
Py_ssize_t *maxvec = NULL;
|
||||||
|
Py_ssize_t *indices = NULL;
|
||||||
|
Py_ssize_t i;
|
||||||
|
|
||||||
|
if (type == &product_type && !_PyArg_NoKeywords("product()", kwds))
|
||||||
|
return NULL;
|
||||||
|
|
||||||
|
assert(PyTuple_Check(args));
|
||||||
|
npools = PyTuple_GET_SIZE(args);
|
||||||
|
|
||||||
|
maxvec = PyMem_Malloc(npools * sizeof(Py_ssize_t));
|
||||||
|
indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
|
||||||
|
if (maxvec == NULL || indices == NULL) {
|
||||||
|
PyErr_NoMemory();
|
||||||
|
goto error;
|
||||||
|
}
|
||||||
|
|
||||||
|
pools = PyTuple_New(npools);
|
||||||
|
if (pools == NULL)
|
||||||
|
goto error;
|
||||||
|
|
||||||
|
for (i=0; i < npools; ++i) {
|
||||||
|
PyObject *item = PyTuple_GET_ITEM(args, i);
|
||||||
|
PyObject *pool = PySequence_Tuple(item);
|
||||||
|
if (pool == NULL)
|
||||||
|
goto error;
|
||||||
|
|
||||||
|
PyTuple_SET_ITEM(pools, i, pool);
|
||||||
|
maxvec[i] = PyTuple_GET_SIZE(pool);
|
||||||
|
indices[i] = 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
/* create productobject structure */
|
||||||
|
lz = (productobject *)type->tp_alloc(type, 0);
|
||||||
|
if (lz == NULL) {
|
||||||
|
Py_DECREF(pools);
|
||||||
|
return NULL;
|
||||||
|
}
|
||||||
|
|
||||||
|
lz->pools = pools;
|
||||||
|
lz->maxvec = maxvec;
|
||||||
|
lz->indices = indices;
|
||||||
|
lz->result = NULL;
|
||||||
|
lz->stopped = 0;
|
||||||
|
|
||||||
|
return (PyObject *)lz;
|
||||||
|
|
||||||
|
error:
|
||||||
|
if (maxvec != NULL)
|
||||||
|
PyMem_Free(maxvec);
|
||||||
|
if (indices != NULL)
|
||||||
|
PyMem_Free(indices);
|
||||||
|
Py_XDECREF(pools);
|
||||||
|
return NULL;
|
||||||
|
}
|
||||||
|
|
||||||
|
static void
|
||||||
|
product_dealloc(productobject *lz)
|
||||||
|
{
|
||||||
|
PyObject_GC_UnTrack(lz);
|
||||||
|
Py_XDECREF(lz->pools);
|
||||||
|
Py_XDECREF(lz->result);
|
||||||
|
PyMem_Free(lz->maxvec);
|
||||||
|
PyMem_Free(lz->indices);
|
||||||
|
Py_TYPE(lz)->tp_free(lz);
|
||||||
|
}
|
||||||
|
|
||||||
|
static int
|
||||||
|
product_traverse(productobject *lz, visitproc visit, void *arg)
|
||||||
|
{
|
||||||
|
Py_VISIT(lz->pools);
|
||||||
|
Py_VISIT(lz->result);
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
static PyObject *
|
||||||
|
product_next(productobject *lz)
|
||||||
|
{
|
||||||
|
PyObject *pool;
|
||||||
|
PyObject *elem;
|
||||||
|
PyObject *tuple_result;
|
||||||
|
PyObject *pools = lz->pools;
|
||||||
|
PyObject *result = lz->result;
|
||||||
|
Py_ssize_t npools = PyTuple_GET_SIZE(pools);
|
||||||
|
Py_ssize_t i;
|
||||||
|
|
||||||
|
if (lz->stopped)
|
||||||
|
return NULL;
|
||||||
|
if (result == NULL) {
|
||||||
|
if (npools == 0)
|
||||||
|
goto empty;
|
||||||
|
result = PyList_New(npools);
|
||||||
|
if (result == NULL)
|
||||||
|
goto empty;
|
||||||
|
lz->result = result;
|
||||||
|
for (i=0; i < npools; i++) {
|
||||||
|
pool = PyTuple_GET_ITEM(pools, i);
|
||||||
|
if (PyTuple_GET_SIZE(pool) == 0)
|
||||||
|
goto empty;
|
||||||
|
elem = PyTuple_GET_ITEM(pool, 0);
|
||||||
|
Py_INCREF(elem);
|
||||||
|
PyList_SET_ITEM(result, i, elem);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Py_ssize_t *indices = lz->indices;
|
||||||
|
Py_ssize_t *maxvec = lz->maxvec;
|
||||||
|
for (i=npools-1 ; i >= 0 ; i--) {
|
||||||
|
pool = PyTuple_GET_ITEM(pools, i);
|
||||||
|
indices[i]++;
|
||||||
|
if (indices[i] == maxvec[i]) {
|
||||||
|
indices[i] = 0;
|
||||||
|
elem = PyTuple_GET_ITEM(pool, 0);
|
||||||
|
Py_INCREF(elem);
|
||||||
|
PyList_SetItem(result, i, elem);
|
||||||
|
} else {
|
||||||
|
elem = PyTuple_GET_ITEM(pool, indices[i]);
|
||||||
|
Py_INCREF(elem);
|
||||||
|
PyList_SetItem(result, i, elem);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (i < 0)
|
||||||
|
return NULL;
|
||||||
|
}
|
||||||
|
|
||||||
|
tuple_result = PySequence_Tuple(result);
|
||||||
|
if (tuple_result == NULL)
|
||||||
|
lz->stopped = 1;
|
||||||
|
return tuple_result;
|
||||||
|
|
||||||
|
empty:
|
||||||
|
lz->stopped = 1;
|
||||||
|
return NULL;
|
||||||
|
}
|
||||||
|
|
||||||
|
PyDoc_STRVAR(product_doc,
|
||||||
|
"product(*iterables) --> product object\n\
|
||||||
|
\n\
|
||||||
|
Cartesian product of input interables. 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\
|
||||||
|
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\
|
||||||
|
on every iteration).\n\n\
|
||||||
|
product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
|
||||||
|
product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
|
||||||
|
|
||||||
|
static PyTypeObject product_type = {
|
||||||
|
PyVarObject_HEAD_INIT(NULL, 0)
|
||||||
|
"itertools.product", /* tp_name */
|
||||||
|
sizeof(productobject), /* tp_basicsize */
|
||||||
|
0, /* tp_itemsize */
|
||||||
|
/* methods */
|
||||||
|
(destructor)product_dealloc, /* tp_dealloc */
|
||||||
|
0, /* tp_print */
|
||||||
|
0, /* tp_getattr */
|
||||||
|
0, /* tp_setattr */
|
||||||
|
0, /* tp_compare */
|
||||||
|
0, /* tp_repr */
|
||||||
|
0, /* tp_as_number */
|
||||||
|
0, /* tp_as_sequence */
|
||||||
|
0, /* tp_as_mapping */
|
||||||
|
0, /* tp_hash */
|
||||||
|
0, /* tp_call */
|
||||||
|
0, /* tp_str */
|
||||||
|
PyObject_GenericGetAttr, /* tp_getattro */
|
||||||
|
0, /* tp_setattro */
|
||||||
|
0, /* tp_as_buffer */
|
||||||
|
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
|
||||||
|
Py_TPFLAGS_BASETYPE, /* tp_flags */
|
||||||
|
product_doc, /* tp_doc */
|
||||||
|
(traverseproc)product_traverse, /* tp_traverse */
|
||||||
|
0, /* tp_clear */
|
||||||
|
0, /* tp_richcompare */
|
||||||
|
0, /* tp_weaklistoffset */
|
||||||
|
PyObject_SelfIter, /* tp_iter */
|
||||||
|
(iternextfunc)product_next, /* tp_iternext */
|
||||||
|
0, /* tp_methods */
|
||||||
|
0, /* tp_members */
|
||||||
|
0, /* tp_getset */
|
||||||
|
0, /* tp_base */
|
||||||
|
0, /* tp_dict */
|
||||||
|
0, /* tp_descr_get */
|
||||||
|
0, /* tp_descr_set */
|
||||||
|
0, /* tp_dictoffset */
|
||||||
|
0, /* tp_init */
|
||||||
|
0, /* tp_alloc */
|
||||||
|
product_new, /* tp_new */
|
||||||
|
PyObject_GC_Del, /* tp_free */
|
||||||
|
};
|
||||||
|
|
||||||
|
|
||||||
/* ifilter object ************************************************************/
|
/* ifilter object ************************************************************/
|
||||||
|
|
||||||
typedef struct {
|
typedef struct {
|
||||||
|
|
@ -2797,6 +3007,7 @@ inititertools(void)
|
||||||
&count_type,
|
&count_type,
|
||||||
&izip_type,
|
&izip_type,
|
||||||
&iziplongest_type,
|
&iziplongest_type,
|
||||||
|
&product_type,
|
||||||
&repeat_type,
|
&repeat_type,
|
||||||
&groupby_type,
|
&groupby_type,
|
||||||
NULL
|
NULL
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue