mirror of
https://github.com/python/cpython.git
synced 2024-12-05 07:43:50 +08:00
826bf0e695
Add repeat, islice, chain, tee, teedataobject, and batched types to module state. Automerge-Triggered-By: GH:erlend-aasland
4721 lines
130 KiB
C
4721 lines
130 KiB
C
#define PY_SSIZE_T_CLEAN
|
|
#include "Python.h"
|
|
#include "pycore_call.h" // _PyObject_CallNoArgs()
|
|
#include "pycore_long.h" // _PyLong_GetZero()
|
|
#include "pycore_moduleobject.h" // _PyModule_GetState()
|
|
#include "pycore_object.h" // _PyObject_GC_TRACK()
|
|
#include "pycore_tuple.h" // _PyTuple_ITEMS()
|
|
#include "structmember.h" // PyMemberDef
|
|
#include <stddef.h> // offsetof()
|
|
|
|
/* Itertools module written and maintained
|
|
by Raymond D. Hettinger <python@rcn.com>
|
|
*/
|
|
|
|
typedef struct {
|
|
PyTypeObject *accumulate_type;
|
|
PyTypeObject *batched_type;
|
|
PyTypeObject *chain_type;
|
|
PyTypeObject *combinations_type;
|
|
PyTypeObject *compress_type;
|
|
PyTypeObject *count_type;
|
|
PyTypeObject *cwr_type;
|
|
PyTypeObject *cycle_type;
|
|
PyTypeObject *dropwhile_type;
|
|
PyTypeObject *filterfalse_type;
|
|
PyTypeObject *groupby_type;
|
|
PyTypeObject *_grouper_type;
|
|
PyTypeObject *islice_type;
|
|
PyTypeObject *pairwise_type;
|
|
PyTypeObject *permutations_type;
|
|
PyTypeObject *product_type;
|
|
PyTypeObject *repeat_type;
|
|
PyTypeObject *starmap_type;
|
|
PyTypeObject *takewhile_type;
|
|
PyTypeObject *tee_type;
|
|
PyTypeObject *teedataobject_type;
|
|
PyTypeObject *ziplongest_type;
|
|
} itertools_state;
|
|
|
|
static inline itertools_state *
|
|
get_module_state(PyObject *mod)
|
|
{
|
|
void *state = _PyModule_GetState(mod);
|
|
assert(state != NULL);
|
|
return (itertools_state *)state;
|
|
}
|
|
|
|
static inline itertools_state *
|
|
get_module_state_by_cls(PyTypeObject *cls)
|
|
{
|
|
void *state = PyType_GetModuleState(cls);
|
|
assert(state != NULL);
|
|
return (itertools_state *)state;
|
|
}
|
|
|
|
static struct PyModuleDef itertoolsmodule;
|
|
|
|
static inline itertools_state *
|
|
find_state_by_type(PyTypeObject *tp)
|
|
{
|
|
PyObject *mod = PyType_GetModuleByDef(tp, &itertoolsmodule);
|
|
assert(mod != NULL);
|
|
return get_module_state(mod);
|
|
}
|
|
|
|
/*[clinic input]
|
|
module itertools
|
|
class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
|
|
class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
|
|
class itertools.teedataobject "teedataobject *" "clinic_state()->teedataobject_type"
|
|
class itertools._tee "teeobject *" "clinic_state()->tee_type"
|
|
class itertools.batched "batchedobject *" "clinic_state()->batched_type"
|
|
class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
|
|
class itertools.dropwhile "dropwhileobject *" "clinic_state()->dropwhile_type"
|
|
class itertools.takewhile "takewhileobject *" "clinic_state()->takewhile_type"
|
|
class itertools.starmap "starmapobject *" "clinic_state()->starmap_type"
|
|
class itertools.chain "chainobject *" "clinic_state()->chain_type"
|
|
class itertools.combinations "combinationsobject *" "clinic_state()->combinations_type"
|
|
class itertools.combinations_with_replacement "cwr_object *" "clinic_state()->cwr_type"
|
|
class itertools.permutations "permutationsobject *" "clinic_state()->permutations_type"
|
|
class itertools.accumulate "accumulateobject *" "clinic_state()->accumulate_type"
|
|
class itertools.compress "compressobject *" "clinic_state()->compress_type"
|
|
class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_type"
|
|
class itertools.count "countobject *" "clinic_state()->count_type"
|
|
class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
|
|
[clinic start generated code]*/
|
|
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=aa48fe4de9d4080f]*/
|
|
|
|
#define clinic_state() (find_state_by_type(type))
|
|
#define clinic_state_by_cls() (get_module_state_by_cls(base_tp))
|
|
#include "clinic/itertoolsmodule.c.h"
|
|
#undef clinic_state_by_cls
|
|
#undef clinic_state
|
|
|
|
/* batched object ************************************************************/
|
|
|
|
/* Note: The built-in zip() function includes a "strict" argument
|
|
that was needed because that function would silently truncate data,
|
|
and there was no easy way for a user to detect the data loss.
|
|
The same reasoning does not apply to batched() which never drops data.
|
|
Instead, batched() produces a shorter tuple which can be handled
|
|
as the user sees fit. If requested, it would be reasonable to add
|
|
"fillvalue" support which had demonstrated value in zip_longest().
|
|
For now, the API is kept simple and clean.
|
|
*/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *it;
|
|
Py_ssize_t batch_size;
|
|
} batchedobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.batched.__new__ as batched_new
|
|
iterable: object
|
|
n: Py_ssize_t
|
|
Batch data into tuples of length n. The last batch may be shorter than n.
|
|
|
|
Loops over the input iterable and accumulates data into tuples
|
|
up to size n. The input is consumed lazily, just enough to
|
|
fill a batch. The result is yielded as soon as a batch is full
|
|
or when the input iterable is exhausted.
|
|
|
|
>>> for batch in batched('ABCDEFG', 3):
|
|
... print(batch)
|
|
...
|
|
('A', 'B', 'C')
|
|
('D', 'E', 'F')
|
|
('G',)
|
|
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
batched_new_impl(PyTypeObject *type, PyObject *iterable, Py_ssize_t n)
|
|
/*[clinic end generated code: output=7ebc954d655371b6 input=ffd70726927c5129]*/
|
|
{
|
|
PyObject *it;
|
|
batchedobject *bo;
|
|
|
|
if (n < 1) {
|
|
/* We could define the n==0 case to return an empty iterator
|
|
but that is at odds with the idea that batching should
|
|
never throw-away input data.
|
|
*/
|
|
PyErr_SetString(PyExc_ValueError, "n must be at least one");
|
|
return NULL;
|
|
}
|
|
it = PyObject_GetIter(iterable);
|
|
if (it == NULL) {
|
|
return NULL;
|
|
}
|
|
|
|
/* create batchedobject structure */
|
|
bo = (batchedobject *)type->tp_alloc(type, 0);
|
|
if (bo == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
bo->batch_size = n;
|
|
bo->it = it;
|
|
return (PyObject *)bo;
|
|
}
|
|
|
|
static void
|
|
batched_dealloc(batchedobject *bo)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(bo);
|
|
PyObject_GC_UnTrack(bo);
|
|
Py_XDECREF(bo->it);
|
|
tp->tp_free(bo);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
batched_traverse(batchedobject *bo, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(bo));
|
|
Py_VISIT(bo->it);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
batched_next(batchedobject *bo)
|
|
{
|
|
Py_ssize_t i;
|
|
Py_ssize_t n = bo->batch_size;
|
|
PyObject *it = bo->it;
|
|
PyObject *item;
|
|
PyObject *result;
|
|
|
|
if (it == NULL) {
|
|
return NULL;
|
|
}
|
|
result = PyTuple_New(n);
|
|
if (result == NULL) {
|
|
return NULL;
|
|
}
|
|
iternextfunc iternext = *Py_TYPE(it)->tp_iternext;
|
|
PyObject **items = _PyTuple_ITEMS(result);
|
|
for (i=0 ; i < n ; i++) {
|
|
item = iternext(it);
|
|
if (item == NULL) {
|
|
goto null_item;
|
|
}
|
|
items[i] = item;
|
|
}
|
|
return result;
|
|
|
|
null_item:
|
|
if (PyErr_Occurred()) {
|
|
if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
|
|
/* Input raised an exception other than StopIteration */
|
|
Py_CLEAR(bo->it);
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
}
|
|
PyErr_Clear();
|
|
}
|
|
if (i == 0) {
|
|
Py_CLEAR(bo->it);
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
}
|
|
_PyTuple_Resize(&result, i);
|
|
return result;
|
|
}
|
|
|
|
static PyType_Slot batched_slots[] = {
|
|
{Py_tp_dealloc, batched_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)batched_new__doc__},
|
|
{Py_tp_traverse, batched_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, batched_next},
|
|
{Py_tp_alloc, PyType_GenericAlloc},
|
|
{Py_tp_new, batched_new},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec batched_spec = {
|
|
.name = "itertools.batched",
|
|
.basicsize = sizeof(batchedobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = batched_slots,
|
|
};
|
|
|
|
|
|
/* pairwise object ***********************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *it;
|
|
PyObject *old;
|
|
} pairwiseobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.pairwise.__new__ as pairwise_new
|
|
iterable: object
|
|
/
|
|
Return an iterator of overlapping pairs taken from the input iterator.
|
|
|
|
s -> (s0,s1), (s1,s2), (s2, s3), ...
|
|
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
|
|
/*[clinic end generated code: output=9f0267062d384456 input=6e7c3cddb431a8d6]*/
|
|
{
|
|
PyObject *it;
|
|
pairwiseobject *po;
|
|
|
|
it = PyObject_GetIter(iterable);
|
|
if (it == NULL) {
|
|
return NULL;
|
|
}
|
|
po = (pairwiseobject *)type->tp_alloc(type, 0);
|
|
if (po == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
po->it = it;
|
|
po->old = NULL;
|
|
return (PyObject *)po;
|
|
}
|
|
|
|
static void
|
|
pairwise_dealloc(pairwiseobject *po)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(po);
|
|
PyObject_GC_UnTrack(po);
|
|
Py_XDECREF(po->it);
|
|
Py_XDECREF(po->old);
|
|
tp->tp_free(po);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
pairwise_traverse(pairwiseobject *po, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(po));
|
|
Py_VISIT(po->it);
|
|
Py_VISIT(po->old);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
pairwise_next(pairwiseobject *po)
|
|
{
|
|
PyObject *it = po->it;
|
|
PyObject *old = po->old;
|
|
PyObject *new, *result;
|
|
|
|
if (it == NULL) {
|
|
return NULL;
|
|
}
|
|
if (old == NULL) {
|
|
po->old = old = (*Py_TYPE(it)->tp_iternext)(it);
|
|
if (old == NULL) {
|
|
Py_CLEAR(po->it);
|
|
return NULL;
|
|
}
|
|
}
|
|
new = (*Py_TYPE(it)->tp_iternext)(it);
|
|
if (new == NULL) {
|
|
Py_CLEAR(po->it);
|
|
Py_CLEAR(po->old);
|
|
return NULL;
|
|
}
|
|
/* Future optimization: Reuse the result tuple as we do in enumerate() */
|
|
result = PyTuple_Pack(2, old, new);
|
|
Py_SETREF(po->old, new);
|
|
return result;
|
|
}
|
|
|
|
static PyType_Slot pairwise_slots[] = {
|
|
{Py_tp_dealloc, pairwise_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)pairwise_new__doc__},
|
|
{Py_tp_traverse, pairwise_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, pairwise_next},
|
|
{Py_tp_alloc, PyType_GenericAlloc},
|
|
{Py_tp_new, pairwise_new},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec pairwise_spec = {
|
|
.name = "itertools.pairwise",
|
|
.basicsize = sizeof(pairwiseobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = pairwise_slots,
|
|
};
|
|
|
|
|
|
/* groupby object ************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *it;
|
|
PyObject *keyfunc;
|
|
PyObject *tgtkey;
|
|
PyObject *currkey;
|
|
PyObject *currvalue;
|
|
const void *currgrouper; /* borrowed reference */
|
|
itertools_state *state;
|
|
} groupbyobject;
|
|
|
|
static PyObject *_grouper_create(groupbyobject *, PyObject *);
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.groupby.__new__
|
|
|
|
iterable as it: object
|
|
Elements to divide into groups according to the key function.
|
|
key as keyfunc: object = None
|
|
A function for computing the group category for each element.
|
|
If the key function is not specified or is None, the element itself
|
|
is used for grouping.
|
|
|
|
make an iterator that returns consecutive keys and groups from the iterable
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
|
|
/*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
|
|
{
|
|
groupbyobject *gbo;
|
|
|
|
gbo = (groupbyobject *)type->tp_alloc(type, 0);
|
|
if (gbo == NULL)
|
|
return NULL;
|
|
gbo->tgtkey = NULL;
|
|
gbo->currkey = NULL;
|
|
gbo->currvalue = NULL;
|
|
gbo->keyfunc = Py_NewRef(keyfunc);
|
|
gbo->it = PyObject_GetIter(it);
|
|
if (gbo->it == NULL) {
|
|
Py_DECREF(gbo);
|
|
return NULL;
|
|
}
|
|
gbo->state = find_state_by_type(type);
|
|
return (PyObject *)gbo;
|
|
}
|
|
|
|
static void
|
|
groupby_dealloc(groupbyobject *gbo)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(gbo);
|
|
PyObject_GC_UnTrack(gbo);
|
|
Py_XDECREF(gbo->it);
|
|
Py_XDECREF(gbo->keyfunc);
|
|
Py_XDECREF(gbo->tgtkey);
|
|
Py_XDECREF(gbo->currkey);
|
|
Py_XDECREF(gbo->currvalue);
|
|
tp->tp_free(gbo);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(gbo));
|
|
Py_VISIT(gbo->it);
|
|
Py_VISIT(gbo->keyfunc);
|
|
Py_VISIT(gbo->tgtkey);
|
|
Py_VISIT(gbo->currkey);
|
|
Py_VISIT(gbo->currvalue);
|
|
return 0;
|
|
}
|
|
|
|
Py_LOCAL_INLINE(int)
|
|
groupby_step(groupbyobject *gbo)
|
|
{
|
|
PyObject *newvalue, *newkey, *oldvalue;
|
|
|
|
newvalue = PyIter_Next(gbo->it);
|
|
if (newvalue == NULL)
|
|
return -1;
|
|
|
|
if (gbo->keyfunc == Py_None) {
|
|
newkey = Py_NewRef(newvalue);
|
|
} else {
|
|
newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
|
|
if (newkey == NULL) {
|
|
Py_DECREF(newvalue);
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
oldvalue = gbo->currvalue;
|
|
gbo->currvalue = newvalue;
|
|
Py_XSETREF(gbo->currkey, newkey);
|
|
Py_XDECREF(oldvalue);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
groupby_next(groupbyobject *gbo)
|
|
{
|
|
PyObject *r, *grouper;
|
|
|
|
gbo->currgrouper = NULL;
|
|
/* skip to next iteration group */
|
|
for (;;) {
|
|
if (gbo->currkey == NULL)
|
|
/* pass */;
|
|
else if (gbo->tgtkey == NULL)
|
|
break;
|
|
else {
|
|
int rcmp;
|
|
|
|
rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
|
|
if (rcmp == -1)
|
|
return NULL;
|
|
else if (rcmp == 0)
|
|
break;
|
|
}
|
|
|
|
if (groupby_step(gbo) < 0)
|
|
return NULL;
|
|
}
|
|
Py_INCREF(gbo->currkey);
|
|
Py_XSETREF(gbo->tgtkey, gbo->currkey);
|
|
|
|
grouper = _grouper_create(gbo, gbo->tgtkey);
|
|
if (grouper == NULL)
|
|
return NULL;
|
|
|
|
r = PyTuple_Pack(2, gbo->currkey, grouper);
|
|
Py_DECREF(grouper);
|
|
return r;
|
|
}
|
|
|
|
static PyObject *
|
|
groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
/* reduce as a 'new' call with an optional 'setstate' if groupby
|
|
* has started
|
|
*/
|
|
PyObject *value;
|
|
if (lz->tgtkey && lz->currkey && lz->currvalue)
|
|
value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
|
|
lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
|
|
else
|
|
value = Py_BuildValue("O(OO)", Py_TYPE(lz),
|
|
lz->it, lz->keyfunc);
|
|
|
|
return value;
|
|
}
|
|
|
|
PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
|
|
|
|
static PyObject *
|
|
groupby_setstate(groupbyobject *lz, PyObject *state)
|
|
{
|
|
PyObject *currkey, *currvalue, *tgtkey;
|
|
if (!PyTuple_Check(state)) {
|
|
PyErr_SetString(PyExc_TypeError, "state is not a tuple");
|
|
return NULL;
|
|
}
|
|
if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
|
|
return NULL;
|
|
}
|
|
Py_INCREF(currkey);
|
|
Py_XSETREF(lz->currkey, currkey);
|
|
Py_INCREF(currvalue);
|
|
Py_XSETREF(lz->currvalue, currvalue);
|
|
Py_INCREF(tgtkey);
|
|
Py_XSETREF(lz->tgtkey, tgtkey);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
|
|
|
|
static PyMethodDef groupby_methods[] = {
|
|
{"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)groupby_setstate, METH_O,
|
|
setstate_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot groupby_slots[] = {
|
|
{Py_tp_dealloc, groupby_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_groupby__doc__},
|
|
{Py_tp_traverse, groupby_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, groupby_next},
|
|
{Py_tp_methods, groupby_methods},
|
|
{Py_tp_new, itertools_groupby},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec groupby_spec = {
|
|
.name = "itertools.groupby",
|
|
.basicsize= sizeof(groupbyobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = groupby_slots,
|
|
};
|
|
|
|
/* _grouper object (internal) ************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *parent;
|
|
PyObject *tgtkey;
|
|
} _grouperobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools._grouper.__new__
|
|
|
|
parent: object(subclass_of='clinic_state_by_cls()->groupby_type')
|
|
tgtkey: object
|
|
/
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
|
|
PyObject *tgtkey)
|
|
/*[clinic end generated code: output=462efb1cdebb5914 input=afe05eb477118f12]*/
|
|
{
|
|
return _grouper_create((groupbyobject*) parent, tgtkey);
|
|
}
|
|
|
|
static PyObject *
|
|
_grouper_create(groupbyobject *parent, PyObject *tgtkey)
|
|
{
|
|
itertools_state *state = parent->state;
|
|
_grouperobject *igo = PyObject_GC_New(_grouperobject, state->_grouper_type);
|
|
if (igo == NULL)
|
|
return NULL;
|
|
igo->parent = Py_NewRef(parent);
|
|
igo->tgtkey = Py_NewRef(tgtkey);
|
|
parent->currgrouper = igo; /* borrowed reference */
|
|
|
|
PyObject_GC_Track(igo);
|
|
return (PyObject *)igo;
|
|
}
|
|
|
|
static void
|
|
_grouper_dealloc(_grouperobject *igo)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(igo);
|
|
PyObject_GC_UnTrack(igo);
|
|
Py_DECREF(igo->parent);
|
|
Py_DECREF(igo->tgtkey);
|
|
PyObject_GC_Del(igo);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(igo));
|
|
Py_VISIT(igo->parent);
|
|
Py_VISIT(igo->tgtkey);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
_grouper_next(_grouperobject *igo)
|
|
{
|
|
groupbyobject *gbo = (groupbyobject *)igo->parent;
|
|
PyObject *r;
|
|
int rcmp;
|
|
|
|
if (gbo->currgrouper != igo)
|
|
return NULL;
|
|
if (gbo->currvalue == NULL) {
|
|
if (groupby_step(gbo) < 0)
|
|
return NULL;
|
|
}
|
|
|
|
assert(gbo->currkey != NULL);
|
|
rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
|
|
if (rcmp <= 0)
|
|
/* got any error or current group is end */
|
|
return NULL;
|
|
|
|
r = gbo->currvalue;
|
|
gbo->currvalue = NULL;
|
|
Py_CLEAR(gbo->currkey);
|
|
|
|
return r;
|
|
}
|
|
|
|
static PyObject *
|
|
_grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
if (((groupbyobject *)lz->parent)->currgrouper != lz) {
|
|
return Py_BuildValue("N(())", _PyEval_GetBuiltin(&_Py_ID(iter)));
|
|
}
|
|
return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
|
|
}
|
|
|
|
static PyMethodDef _grouper_methods[] = {
|
|
{"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot _grouper_slots[] = {
|
|
{Py_tp_dealloc, _grouper_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_traverse, _grouper_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, _grouper_next},
|
|
{Py_tp_methods, _grouper_methods},
|
|
{Py_tp_new, itertools__grouper},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec _grouper_spec = {
|
|
.name = "itertools._grouper",
|
|
.basicsize = sizeof(_grouperobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = _grouper_slots,
|
|
};
|
|
|
|
|
|
/* tee object and with supporting function and objects ***********************/
|
|
|
|
/* The teedataobject pre-allocates space for LINKCELLS number of objects.
|
|
To help the object fit neatly inside cache lines (space for 16 to 32
|
|
pointers), the value should be a multiple of 16 minus space for
|
|
the other structure members including PyHEAD overhead. The larger the
|
|
value, the less memory overhead per object and the less time spent
|
|
allocating/deallocating new links. The smaller the number, the less
|
|
wasted space and the more rapid freeing of older data.
|
|
*/
|
|
#define LINKCELLS 57
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *it;
|
|
int numread; /* 0 <= numread <= LINKCELLS */
|
|
int running;
|
|
PyObject *nextlink;
|
|
PyObject *(values[LINKCELLS]);
|
|
} teedataobject;
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
teedataobject *dataobj;
|
|
int index; /* 0 <= index <= LINKCELLS */
|
|
PyObject *weakreflist;
|
|
itertools_state *state;
|
|
} teeobject;
|
|
|
|
static PyObject *
|
|
teedataobject_newinternal(itertools_state *state, PyObject *it)
|
|
{
|
|
teedataobject *tdo;
|
|
|
|
tdo = PyObject_GC_New(teedataobject, state->teedataobject_type);
|
|
if (tdo == NULL)
|
|
return NULL;
|
|
|
|
tdo->running = 0;
|
|
tdo->numread = 0;
|
|
tdo->nextlink = NULL;
|
|
tdo->it = Py_NewRef(it);
|
|
PyObject_GC_Track(tdo);
|
|
return (PyObject *)tdo;
|
|
}
|
|
|
|
static PyObject *
|
|
teedataobject_jumplink(itertools_state *state, teedataobject *tdo)
|
|
{
|
|
if (tdo->nextlink == NULL)
|
|
tdo->nextlink = teedataobject_newinternal(state, tdo->it);
|
|
return Py_XNewRef(tdo->nextlink);
|
|
}
|
|
|
|
static PyObject *
|
|
teedataobject_getitem(teedataobject *tdo, int i)
|
|
{
|
|
PyObject *value;
|
|
|
|
assert(i < LINKCELLS);
|
|
if (i < tdo->numread)
|
|
value = tdo->values[i];
|
|
else {
|
|
/* this is the lead iterator, so fetch more data */
|
|
assert(i == tdo->numread);
|
|
if (tdo->running) {
|
|
PyErr_SetString(PyExc_RuntimeError,
|
|
"cannot re-enter the tee iterator");
|
|
return NULL;
|
|
}
|
|
tdo->running = 1;
|
|
value = PyIter_Next(tdo->it);
|
|
tdo->running = 0;
|
|
if (value == NULL)
|
|
return NULL;
|
|
tdo->numread++;
|
|
tdo->values[i] = value;
|
|
}
|
|
return Py_NewRef(value);
|
|
}
|
|
|
|
static int
|
|
teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
|
|
{
|
|
int i;
|
|
|
|
Py_VISIT(Py_TYPE(tdo));
|
|
Py_VISIT(tdo->it);
|
|
for (i = 0; i < tdo->numread; i++)
|
|
Py_VISIT(tdo->values[i]);
|
|
Py_VISIT(tdo->nextlink);
|
|
return 0;
|
|
}
|
|
|
|
static void
|
|
teedataobject_safe_decref(PyObject *obj, PyTypeObject *tdo_type)
|
|
{
|
|
while (obj && Py_IS_TYPE(obj, tdo_type) &&
|
|
Py_REFCNT(obj) == 1) {
|
|
PyObject *nextlink = ((teedataobject *)obj)->nextlink;
|
|
((teedataobject *)obj)->nextlink = NULL;
|
|
Py_SETREF(obj, nextlink);
|
|
}
|
|
Py_XDECREF(obj);
|
|
}
|
|
|
|
static int
|
|
teedataobject_clear(teedataobject *tdo)
|
|
{
|
|
int i;
|
|
PyObject *tmp;
|
|
|
|
Py_CLEAR(tdo->it);
|
|
for (i=0 ; i<tdo->numread ; i++)
|
|
Py_CLEAR(tdo->values[i]);
|
|
tmp = tdo->nextlink;
|
|
tdo->nextlink = NULL;
|
|
itertools_state *state = get_module_state_by_cls(Py_TYPE(tdo));
|
|
teedataobject_safe_decref(tmp, state->teedataobject_type);
|
|
return 0;
|
|
}
|
|
|
|
static void
|
|
teedataobject_dealloc(teedataobject *tdo)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(tdo);
|
|
PyObject_GC_UnTrack(tdo);
|
|
teedataobject_clear(tdo);
|
|
PyObject_GC_Del(tdo);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static PyObject *
|
|
teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
int i;
|
|
/* create a temporary list of already iterated values */
|
|
PyObject *values = PyList_New(tdo->numread);
|
|
|
|
if (!values)
|
|
return NULL;
|
|
for (i=0 ; i<tdo->numread ; i++) {
|
|
Py_INCREF(tdo->values[i]);
|
|
PyList_SET_ITEM(values, i, tdo->values[i]);
|
|
}
|
|
return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
|
|
values,
|
|
tdo->nextlink ? tdo->nextlink : Py_None);
|
|
}
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.teedataobject.__new__
|
|
iterable as it: object
|
|
values: object(subclass_of='&PyList_Type')
|
|
next: object
|
|
/
|
|
Data container common to multiple tee objects.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
|
|
PyObject *values, PyObject *next)
|
|
/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
|
|
{
|
|
teedataobject *tdo;
|
|
Py_ssize_t i, len;
|
|
|
|
itertools_state *state = get_module_state_by_cls(type);
|
|
assert(type == state->teedataobject_type);
|
|
|
|
tdo = (teedataobject *)teedataobject_newinternal(state, it);
|
|
if (!tdo)
|
|
return NULL;
|
|
|
|
len = PyList_GET_SIZE(values);
|
|
if (len > LINKCELLS)
|
|
goto err;
|
|
for (i=0; i<len; i++) {
|
|
tdo->values[i] = PyList_GET_ITEM(values, i);
|
|
Py_INCREF(tdo->values[i]);
|
|
}
|
|
/* len <= LINKCELLS < INT_MAX */
|
|
tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
|
|
|
|
if (len == LINKCELLS) {
|
|
if (next != Py_None) {
|
|
if (!Py_IS_TYPE(next, state->teedataobject_type))
|
|
goto err;
|
|
assert(tdo->nextlink == NULL);
|
|
tdo->nextlink = Py_NewRef(next);
|
|
}
|
|
} else {
|
|
if (next != Py_None)
|
|
goto err; /* shouldn't have a next if we are not full */
|
|
}
|
|
return (PyObject*)tdo;
|
|
|
|
err:
|
|
Py_XDECREF(tdo);
|
|
PyErr_SetString(PyExc_ValueError, "Invalid arguments");
|
|
return NULL;
|
|
}
|
|
|
|
static PyMethodDef teedataobject_methods[] = {
|
|
{"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot teedataobject_slots[] = {
|
|
{Py_tp_dealloc, teedataobject_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_teedataobject__doc__},
|
|
{Py_tp_traverse, teedataobject_traverse},
|
|
{Py_tp_clear, teedataobject_clear},
|
|
{Py_tp_methods, teedataobject_methods},
|
|
{Py_tp_new, itertools_teedataobject},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec teedataobject_spec = {
|
|
.name = "itertools._tee_dataobject",
|
|
.basicsize = sizeof(teedataobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = teedataobject_slots,
|
|
};
|
|
|
|
|
|
static PyObject *
|
|
tee_next(teeobject *to)
|
|
{
|
|
PyObject *value, *link;
|
|
|
|
if (to->index >= LINKCELLS) {
|
|
link = teedataobject_jumplink(to->state, to->dataobj);
|
|
if (link == NULL)
|
|
return NULL;
|
|
Py_SETREF(to->dataobj, (teedataobject *)link);
|
|
to->index = 0;
|
|
}
|
|
value = teedataobject_getitem(to->dataobj, to->index);
|
|
if (value == NULL)
|
|
return NULL;
|
|
to->index++;
|
|
return value;
|
|
}
|
|
|
|
static int
|
|
tee_traverse(teeobject *to, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(to));
|
|
Py_VISIT((PyObject *)to->dataobj);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
teeobject *newto;
|
|
|
|
newto = PyObject_GC_New(teeobject, Py_TYPE(to));
|
|
if (newto == NULL)
|
|
return NULL;
|
|
newto->dataobj = (teedataobject*)Py_NewRef(to->dataobj);
|
|
newto->index = to->index;
|
|
newto->weakreflist = NULL;
|
|
newto->state = to->state;
|
|
PyObject_GC_Track(newto);
|
|
return (PyObject *)newto;
|
|
}
|
|
|
|
PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
|
|
|
|
static PyObject *
|
|
tee_fromiterable(itertools_state *state, PyObject *iterable)
|
|
{
|
|
teeobject *to;
|
|
PyObject *it;
|
|
|
|
it = PyObject_GetIter(iterable);
|
|
if (it == NULL)
|
|
return NULL;
|
|
if (PyObject_TypeCheck(it, state->tee_type)) {
|
|
to = (teeobject *)tee_copy((teeobject *)it, NULL);
|
|
goto done;
|
|
}
|
|
|
|
PyObject *dataobj = teedataobject_newinternal(state, it);
|
|
if (!dataobj) {
|
|
to = NULL;
|
|
goto done;
|
|
}
|
|
to = PyObject_GC_New(teeobject, state->tee_type);
|
|
if (to == NULL) {
|
|
Py_DECREF(dataobj);
|
|
goto done;
|
|
}
|
|
to->dataobj = (teedataobject *)dataobj;
|
|
to->index = 0;
|
|
to->weakreflist = NULL;
|
|
to->state = state;
|
|
PyObject_GC_Track(to);
|
|
done:
|
|
Py_DECREF(it);
|
|
return (PyObject *)to;
|
|
}
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools._tee.__new__
|
|
iterable: object
|
|
/
|
|
Iterator wrapped to make it copyable.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
|
|
/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
|
|
{
|
|
itertools_state *state = get_module_state_by_cls(type);
|
|
return tee_fromiterable(state, iterable);
|
|
}
|
|
|
|
static int
|
|
tee_clear(teeobject *to)
|
|
{
|
|
if (to->weakreflist != NULL)
|
|
PyObject_ClearWeakRefs((PyObject *) to);
|
|
Py_CLEAR(to->dataobj);
|
|
return 0;
|
|
}
|
|
|
|
static void
|
|
tee_dealloc(teeobject *to)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(to);
|
|
PyObject_GC_UnTrack(to);
|
|
tee_clear(to);
|
|
PyObject_GC_Del(to);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static PyObject *
|
|
tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
|
|
}
|
|
|
|
static PyObject *
|
|
tee_setstate(teeobject *to, PyObject *state)
|
|
{
|
|
teedataobject *tdo;
|
|
int index;
|
|
if (!PyTuple_Check(state)) {
|
|
PyErr_SetString(PyExc_TypeError, "state is not a tuple");
|
|
return NULL;
|
|
}
|
|
PyTypeObject *tdo_type = to->state->teedataobject_type;
|
|
if (!PyArg_ParseTuple(state, "O!i", tdo_type, &tdo, &index)) {
|
|
return NULL;
|
|
}
|
|
if (index < 0 || index > LINKCELLS) {
|
|
PyErr_SetString(PyExc_ValueError, "Index out of range");
|
|
return NULL;
|
|
}
|
|
Py_INCREF(tdo);
|
|
Py_XSETREF(to->dataobj, tdo);
|
|
to->index = index;
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef tee_methods[] = {
|
|
{"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
|
|
{"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
|
|
{"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyMemberDef tee_members[] = {
|
|
{"__weaklistoffset__", T_PYSSIZET, offsetof(teeobject, weakreflist), READONLY},
|
|
{NULL},
|
|
};
|
|
|
|
static PyType_Slot tee_slots[] = {
|
|
{Py_tp_dealloc, tee_dealloc},
|
|
{Py_tp_doc, (void *)itertools__tee__doc__},
|
|
{Py_tp_traverse, tee_traverse},
|
|
{Py_tp_clear, tee_clear},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, tee_next},
|
|
{Py_tp_methods, tee_methods},
|
|
{Py_tp_members, tee_members},
|
|
{Py_tp_new, itertools__tee},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec tee_spec = {
|
|
.name = "itertools._tee",
|
|
.basicsize = sizeof(teeobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = tee_slots,
|
|
};
|
|
|
|
/*[clinic input]
|
|
itertools.tee
|
|
iterable: object
|
|
n: Py_ssize_t = 2
|
|
/
|
|
Returns a tuple of n independent iterators.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
|
|
/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
|
|
{
|
|
Py_ssize_t i;
|
|
PyObject *it, *copyable, *copyfunc, *result;
|
|
|
|
if (n < 0) {
|
|
PyErr_SetString(PyExc_ValueError, "n must be >= 0");
|
|
return NULL;
|
|
}
|
|
result = PyTuple_New(n);
|
|
if (result == NULL)
|
|
return NULL;
|
|
if (n == 0)
|
|
return result;
|
|
it = PyObject_GetIter(iterable);
|
|
if (it == NULL) {
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
}
|
|
|
|
if (_PyObject_LookupAttr(it, &_Py_ID(__copy__), ©func) < 0) {
|
|
Py_DECREF(it);
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
}
|
|
if (copyfunc != NULL) {
|
|
copyable = it;
|
|
}
|
|
else {
|
|
itertools_state *state = get_module_state(module);
|
|
copyable = tee_fromiterable(state, it);
|
|
Py_DECREF(it);
|
|
if (copyable == NULL) {
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
}
|
|
copyfunc = PyObject_GetAttr(copyable, &_Py_ID(__copy__));
|
|
if (copyfunc == NULL) {
|
|
Py_DECREF(copyable);
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
PyTuple_SET_ITEM(result, 0, copyable);
|
|
for (i = 1; i < n; i++) {
|
|
copyable = _PyObject_CallNoArgs(copyfunc);
|
|
if (copyable == NULL) {
|
|
Py_DECREF(copyfunc);
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
}
|
|
PyTuple_SET_ITEM(result, i, copyable);
|
|
}
|
|
Py_DECREF(copyfunc);
|
|
return result;
|
|
}
|
|
|
|
|
|
/* cycle object **************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *it;
|
|
PyObject *saved;
|
|
Py_ssize_t index;
|
|
int firstpass;
|
|
} cycleobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.cycle.__new__
|
|
iterable: object
|
|
/
|
|
Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
|
|
/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
|
|
{
|
|
PyObject *it;
|
|
PyObject *saved;
|
|
cycleobject *lz;
|
|
|
|
/* Get iterator. */
|
|
it = PyObject_GetIter(iterable);
|
|
if (it == NULL)
|
|
return NULL;
|
|
|
|
saved = PyList_New(0);
|
|
if (saved == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
|
|
/* create cycleobject structure */
|
|
lz = (cycleobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(it);
|
|
Py_DECREF(saved);
|
|
return NULL;
|
|
}
|
|
lz->it = it;
|
|
lz->saved = saved;
|
|
lz->index = 0;
|
|
lz->firstpass = 0;
|
|
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
cycle_dealloc(cycleobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->it);
|
|
Py_XDECREF(lz->saved);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->it);
|
|
Py_VISIT(lz->saved);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
cycle_next(cycleobject *lz)
|
|
{
|
|
PyObject *item;
|
|
|
|
if (lz->it != NULL) {
|
|
item = PyIter_Next(lz->it);
|
|
if (item != NULL) {
|
|
if (lz->firstpass)
|
|
return item;
|
|
if (PyList_Append(lz->saved, item)) {
|
|
Py_DECREF(item);
|
|
return NULL;
|
|
}
|
|
return item;
|
|
}
|
|
/* Note: StopIteration is already cleared by PyIter_Next() */
|
|
if (PyErr_Occurred())
|
|
return NULL;
|
|
Py_CLEAR(lz->it);
|
|
}
|
|
if (PyList_GET_SIZE(lz->saved) == 0)
|
|
return NULL;
|
|
item = PyList_GET_ITEM(lz->saved, lz->index);
|
|
lz->index++;
|
|
if (lz->index >= PyList_GET_SIZE(lz->saved))
|
|
lz->index = 0;
|
|
return Py_NewRef(item);
|
|
}
|
|
|
|
static PyObject *
|
|
cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
/* Create a new cycle with the iterator tuple, then set the saved state */
|
|
if (lz->it == NULL) {
|
|
PyObject *it = PyObject_GetIter(lz->saved);
|
|
if (it == NULL)
|
|
return NULL;
|
|
if (lz->index != 0) {
|
|
PyObject *res = _PyObject_CallMethod(it, &_Py_ID(__setstate__),
|
|
"n", lz->index);
|
|
if (res == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
Py_DECREF(res);
|
|
}
|
|
return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
|
|
}
|
|
return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
|
|
lz->firstpass ? Py_True : Py_False);
|
|
}
|
|
|
|
static PyObject *
|
|
cycle_setstate(cycleobject *lz, PyObject *state)
|
|
{
|
|
PyObject *saved=NULL;
|
|
int firstpass;
|
|
if (!PyTuple_Check(state)) {
|
|
PyErr_SetString(PyExc_TypeError, "state is not a tuple");
|
|
return NULL;
|
|
}
|
|
// The second item can be 1/0 in old pickles and True/False in new pickles
|
|
if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
|
|
return NULL;
|
|
}
|
|
Py_INCREF(saved);
|
|
Py_XSETREF(lz->saved, saved);
|
|
lz->firstpass = firstpass != 0;
|
|
lz->index = 0;
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef cycle_methods[] = {
|
|
{"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)cycle_setstate, METH_O,
|
|
setstate_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot cycle_slots[] = {
|
|
{Py_tp_dealloc, cycle_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_cycle__doc__},
|
|
{Py_tp_traverse, cycle_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, cycle_next},
|
|
{Py_tp_methods, cycle_methods},
|
|
{Py_tp_new, itertools_cycle},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec cycle_spec = {
|
|
.name = "itertools.cycle",
|
|
.basicsize = sizeof(cycleobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = cycle_slots,
|
|
};
|
|
|
|
|
|
/* dropwhile object **********************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *func;
|
|
PyObject *it;
|
|
long start;
|
|
} dropwhileobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.dropwhile.__new__
|
|
predicate as func: object
|
|
iterable as seq: object
|
|
/
|
|
Drop items from the iterable while predicate(item) is true.
|
|
|
|
Afterwards, return every element until the iterable is exhausted.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
|
|
/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
|
|
{
|
|
PyObject *it;
|
|
dropwhileobject *lz;
|
|
|
|
/* Get iterator. */
|
|
it = PyObject_GetIter(seq);
|
|
if (it == NULL)
|
|
return NULL;
|
|
|
|
/* create dropwhileobject structure */
|
|
lz = (dropwhileobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
lz->func = Py_NewRef(func);
|
|
lz->it = it;
|
|
lz->start = 0;
|
|
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
dropwhile_dealloc(dropwhileobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->func);
|
|
Py_XDECREF(lz->it);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->it);
|
|
Py_VISIT(lz->func);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
dropwhile_next(dropwhileobject *lz)
|
|
{
|
|
PyObject *item, *good;
|
|
PyObject *it = lz->it;
|
|
long ok;
|
|
PyObject *(*iternext)(PyObject *);
|
|
|
|
iternext = *Py_TYPE(it)->tp_iternext;
|
|
for (;;) {
|
|
item = iternext(it);
|
|
if (item == NULL)
|
|
return NULL;
|
|
if (lz->start == 1)
|
|
return item;
|
|
|
|
good = PyObject_CallOneArg(lz->func, item);
|
|
if (good == NULL) {
|
|
Py_DECREF(item);
|
|
return NULL;
|
|
}
|
|
ok = PyObject_IsTrue(good);
|
|
Py_DECREF(good);
|
|
if (ok == 0) {
|
|
lz->start = 1;
|
|
return item;
|
|
}
|
|
Py_DECREF(item);
|
|
if (ok < 0)
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
static PyObject *
|
|
dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
|
|
}
|
|
|
|
static PyObject *
|
|
dropwhile_setstate(dropwhileobject *lz, PyObject *state)
|
|
{
|
|
int start = PyObject_IsTrue(state);
|
|
if (start < 0)
|
|
return NULL;
|
|
lz->start = start;
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef dropwhile_methods[] = {
|
|
{"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
|
|
setstate_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot dropwhile_slots[] = {
|
|
{Py_tp_dealloc, dropwhile_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_dropwhile__doc__},
|
|
{Py_tp_traverse, dropwhile_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, dropwhile_next},
|
|
{Py_tp_methods, dropwhile_methods},
|
|
{Py_tp_new, itertools_dropwhile},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec dropwhile_spec = {
|
|
.name = "itertools.dropwhile",
|
|
.basicsize = sizeof(dropwhileobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = dropwhile_slots,
|
|
};
|
|
|
|
|
|
/* takewhile object **********************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *func;
|
|
PyObject *it;
|
|
long stop;
|
|
} takewhileobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.takewhile.__new__
|
|
predicate as func: object
|
|
iterable as seq: object
|
|
/
|
|
Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
|
|
/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
|
|
{
|
|
PyObject *it;
|
|
takewhileobject *lz;
|
|
|
|
/* Get iterator. */
|
|
it = PyObject_GetIter(seq);
|
|
if (it == NULL)
|
|
return NULL;
|
|
|
|
/* create takewhileobject structure */
|
|
lz = (takewhileobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
lz->func = Py_NewRef(func);
|
|
lz->it = it;
|
|
lz->stop = 0;
|
|
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
takewhile_dealloc(takewhileobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->func);
|
|
Py_XDECREF(lz->it);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->it);
|
|
Py_VISIT(lz->func);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
takewhile_next(takewhileobject *lz)
|
|
{
|
|
PyObject *item, *good;
|
|
PyObject *it = lz->it;
|
|
long ok;
|
|
|
|
if (lz->stop == 1)
|
|
return NULL;
|
|
|
|
item = (*Py_TYPE(it)->tp_iternext)(it);
|
|
if (item == NULL)
|
|
return NULL;
|
|
|
|
good = PyObject_CallOneArg(lz->func, item);
|
|
if (good == NULL) {
|
|
Py_DECREF(item);
|
|
return NULL;
|
|
}
|
|
ok = PyObject_IsTrue(good);
|
|
Py_DECREF(good);
|
|
if (ok > 0)
|
|
return item;
|
|
Py_DECREF(item);
|
|
if (ok == 0)
|
|
lz->stop = 1;
|
|
return NULL;
|
|
}
|
|
|
|
static PyObject *
|
|
takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
|
|
}
|
|
|
|
static PyObject *
|
|
takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
|
|
{
|
|
int stop = PyObject_IsTrue(state);
|
|
|
|
if (stop < 0)
|
|
return NULL;
|
|
lz->stop = stop;
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef takewhile_reduce_methods[] = {
|
|
{"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
|
|
setstate_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot takewhile_slots[] = {
|
|
{Py_tp_dealloc, takewhile_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_takewhile__doc__},
|
|
{Py_tp_traverse, takewhile_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, takewhile_next},
|
|
{Py_tp_methods, takewhile_reduce_methods},
|
|
{Py_tp_new, itertools_takewhile},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec takewhile_spec = {
|
|
.name = "itertools.takewhile",
|
|
.basicsize = sizeof(takewhileobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = takewhile_slots,
|
|
};
|
|
|
|
|
|
/* islice object *************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *it;
|
|
Py_ssize_t next;
|
|
Py_ssize_t stop;
|
|
Py_ssize_t step;
|
|
Py_ssize_t cnt;
|
|
} isliceobject;
|
|
|
|
static PyObject *
|
|
islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
|
{
|
|
PyObject *seq;
|
|
Py_ssize_t start=0, stop=-1, step=1;
|
|
PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
|
|
Py_ssize_t numargs;
|
|
isliceobject *lz;
|
|
|
|
itertools_state *st = find_state_by_type(type);
|
|
PyTypeObject *islice_type = st->islice_type;
|
|
if ((type == islice_type || type->tp_init == islice_type->tp_init) &&
|
|
!_PyArg_NoKeywords("islice", kwds))
|
|
return NULL;
|
|
|
|
if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
|
|
return NULL;
|
|
|
|
numargs = PyTuple_Size(args);
|
|
if (numargs == 2) {
|
|
if (a1 != Py_None) {
|
|
stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
|
|
if (stop == -1) {
|
|
if (PyErr_Occurred())
|
|
PyErr_Clear();
|
|
PyErr_SetString(PyExc_ValueError,
|
|
"Stop argument for islice() must be None or "
|
|
"an integer: 0 <= x <= sys.maxsize.");
|
|
return NULL;
|
|
}
|
|
}
|
|
} else {
|
|
if (a1 != Py_None)
|
|
start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
|
|
if (start == -1 && PyErr_Occurred())
|
|
PyErr_Clear();
|
|
if (a2 != Py_None) {
|
|
stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
|
|
if (stop == -1) {
|
|
if (PyErr_Occurred())
|
|
PyErr_Clear();
|
|
PyErr_SetString(PyExc_ValueError,
|
|
"Stop argument for islice() must be None or "
|
|
"an integer: 0 <= x <= sys.maxsize.");
|
|
return NULL;
|
|
}
|
|
}
|
|
}
|
|
if (start<0 || stop<-1) {
|
|
PyErr_SetString(PyExc_ValueError,
|
|
"Indices for islice() must be None or "
|
|
"an integer: 0 <= x <= sys.maxsize.");
|
|
return NULL;
|
|
}
|
|
|
|
if (a3 != NULL) {
|
|
if (a3 != Py_None)
|
|
step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
|
|
if (step == -1 && PyErr_Occurred())
|
|
PyErr_Clear();
|
|
}
|
|
if (step<1) {
|
|
PyErr_SetString(PyExc_ValueError,
|
|
"Step for islice() must be a positive integer or None.");
|
|
return NULL;
|
|
}
|
|
|
|
/* Get iterator. */
|
|
it = PyObject_GetIter(seq);
|
|
if (it == NULL)
|
|
return NULL;
|
|
|
|
/* create isliceobject structure */
|
|
lz = (isliceobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
lz->it = it;
|
|
lz->next = start;
|
|
lz->stop = stop;
|
|
lz->step = step;
|
|
lz->cnt = 0L;
|
|
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
islice_dealloc(isliceobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->it);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
islice_traverse(isliceobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->it);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
islice_next(isliceobject *lz)
|
|
{
|
|
PyObject *item;
|
|
PyObject *it = lz->it;
|
|
Py_ssize_t stop = lz->stop;
|
|
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)
|
|
goto empty;
|
|
Py_DECREF(item);
|
|
lz->cnt++;
|
|
}
|
|
if (stop != -1 && lz->cnt >= stop)
|
|
goto empty;
|
|
item = iternext(it);
|
|
if (item == NULL)
|
|
goto empty;
|
|
lz->cnt++;
|
|
oldnext = lz->next;
|
|
/* The (size_t) cast below avoids the danger of undefined
|
|
behaviour from signed integer overflow. */
|
|
lz->next += (size_t)lz->step;
|
|
if (lz->next < oldnext || (stop != -1 && lz->next > stop))
|
|
lz->next = stop;
|
|
return item;
|
|
|
|
empty:
|
|
Py_CLEAR(lz->it);
|
|
return NULL;
|
|
}
|
|
|
|
static PyObject *
|
|
islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
/* When unpickled, generate a new object with the same bounds,
|
|
* 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_NewRef(Py_None);
|
|
} else {
|
|
stop = PyLong_FromSsize_t(lz->stop);
|
|
if (stop == NULL)
|
|
return NULL;
|
|
}
|
|
return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
|
|
lz->it, lz->next, stop, lz->step,
|
|
lz->cnt);
|
|
}
|
|
|
|
static PyObject *
|
|
islice_setstate(isliceobject *lz, PyObject *state)
|
|
{
|
|
Py_ssize_t cnt = PyLong_AsSsize_t(state);
|
|
|
|
if (cnt == -1 && PyErr_Occurred())
|
|
return NULL;
|
|
lz->cnt = cnt;
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef islice_methods[] = {
|
|
{"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)islice_setstate, METH_O,
|
|
setstate_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
PyDoc_STRVAR(islice_doc,
|
|
"islice(iterable, stop) --> islice object\n\
|
|
islice(iterable, start, stop[, step]) --> islice object\n\
|
|
\n\
|
|
Return an iterator whose next() method returns selected values from an\n\
|
|
iterable. If start is specified, will skip all preceding elements;\n\
|
|
otherwise, start defaults to zero. Step defaults to one. If\n\
|
|
specified as another value, step determines how many values are\n\
|
|
skipped between successive calls. Works like a slice() on a list\n\
|
|
but returns an iterator.");
|
|
|
|
static PyType_Slot islice_slots[] = {
|
|
{Py_tp_dealloc, islice_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)islice_doc},
|
|
{Py_tp_traverse, islice_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, islice_next},
|
|
{Py_tp_methods, islice_methods},
|
|
{Py_tp_new, islice_new},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec islice_spec = {
|
|
.name = "itertools.islice",
|
|
.basicsize = sizeof(isliceobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = islice_slots,
|
|
};
|
|
|
|
|
|
/* starmap object ************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *func;
|
|
PyObject *it;
|
|
} starmapobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.starmap.__new__
|
|
function as func: object
|
|
iterable as seq: object
|
|
/
|
|
Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
|
|
/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
|
|
{
|
|
PyObject *it;
|
|
starmapobject *lz;
|
|
|
|
/* Get iterator. */
|
|
it = PyObject_GetIter(seq);
|
|
if (it == NULL)
|
|
return NULL;
|
|
|
|
/* create starmapobject structure */
|
|
lz = (starmapobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
lz->func = Py_NewRef(func);
|
|
lz->it = it;
|
|
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
starmap_dealloc(starmapobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->func);
|
|
Py_XDECREF(lz->it);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->it);
|
|
Py_VISIT(lz->func);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
starmap_next(starmapobject *lz)
|
|
{
|
|
PyObject *args;
|
|
PyObject *result;
|
|
PyObject *it = lz->it;
|
|
|
|
args = (*Py_TYPE(it)->tp_iternext)(it);
|
|
if (args == NULL)
|
|
return NULL;
|
|
if (!PyTuple_CheckExact(args)) {
|
|
PyObject *newargs = PySequence_Tuple(args);
|
|
Py_DECREF(args);
|
|
if (newargs == NULL)
|
|
return NULL;
|
|
args = newargs;
|
|
}
|
|
result = PyObject_Call(lz->func, args, NULL);
|
|
Py_DECREF(args);
|
|
return result;
|
|
}
|
|
|
|
static PyObject *
|
|
starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
/* Just pickle the iterator */
|
|
return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
|
|
}
|
|
|
|
static PyMethodDef starmap_methods[] = {
|
|
{"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot starmap_slots[] = {
|
|
{Py_tp_dealloc, starmap_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_starmap__doc__},
|
|
{Py_tp_traverse, starmap_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, starmap_next},
|
|
{Py_tp_methods, starmap_methods},
|
|
{Py_tp_new, itertools_starmap},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec starmap_spec = {
|
|
.name = "itertools.starmap",
|
|
.basicsize = sizeof(starmapobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = starmap_slots,
|
|
};
|
|
|
|
|
|
/* chain object **************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *source; /* Iterator over input iterables */
|
|
PyObject *active; /* Currently running input iterator */
|
|
} chainobject;
|
|
|
|
static PyObject *
|
|
chain_new_internal(PyTypeObject *type, PyObject *source)
|
|
{
|
|
chainobject *lz;
|
|
|
|
lz = (chainobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(source);
|
|
return NULL;
|
|
}
|
|
|
|
lz->source = source;
|
|
lz->active = NULL;
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static PyObject *
|
|
chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
|
{
|
|
PyObject *source;
|
|
|
|
itertools_state *state = find_state_by_type(type);
|
|
PyTypeObject *chain_type = state->chain_type;
|
|
if ((type == chain_type || type->tp_init == chain_type->tp_init) &&
|
|
!_PyArg_NoKeywords("chain", kwds))
|
|
return NULL;
|
|
|
|
source = PyObject_GetIter(args);
|
|
if (source == NULL)
|
|
return NULL;
|
|
|
|
return chain_new_internal(type, source);
|
|
}
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.chain.from_iterable
|
|
iterable as arg: object
|
|
/
|
|
Alternative chain() constructor taking a single iterable argument that evaluates lazily.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
|
|
/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
|
|
{
|
|
PyObject *source;
|
|
|
|
source = PyObject_GetIter(arg);
|
|
if (source == NULL)
|
|
return NULL;
|
|
|
|
return chain_new_internal(type, source);
|
|
}
|
|
|
|
static void
|
|
chain_dealloc(chainobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->active);
|
|
Py_XDECREF(lz->source);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
chain_traverse(chainobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->source);
|
|
Py_VISIT(lz->active);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
chain_next(chainobject *lz)
|
|
{
|
|
PyObject *item;
|
|
|
|
/* lz->source is the iterator of iterables. If it's NULL, we've already
|
|
* consumed them all. lz->active is the current iterator. If it's NULL,
|
|
* we should grab a new one from lz->source. */
|
|
while (lz->source != NULL) {
|
|
if (lz->active == NULL) {
|
|
PyObject *iterable = PyIter_Next(lz->source);
|
|
if (iterable == NULL) {
|
|
Py_CLEAR(lz->source);
|
|
return NULL; /* no more input sources */
|
|
}
|
|
lz->active = PyObject_GetIter(iterable);
|
|
Py_DECREF(iterable);
|
|
if (lz->active == NULL) {
|
|
Py_CLEAR(lz->source);
|
|
return NULL; /* input not iterable */
|
|
}
|
|
}
|
|
item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
|
|
if (item != NULL)
|
|
return item;
|
|
if (PyErr_Occurred()) {
|
|
if (PyErr_ExceptionMatches(PyExc_StopIteration))
|
|
PyErr_Clear();
|
|
else
|
|
return NULL; /* input raised an exception */
|
|
}
|
|
/* lz->active is consumed, try with the next iterable. */
|
|
Py_CLEAR(lz->active);
|
|
}
|
|
/* Everything had been consumed already. */
|
|
return NULL;
|
|
}
|
|
|
|
static PyObject *
|
|
chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
if (lz->source) {
|
|
/* we can't pickle function objects (itertools.from_iterable) so
|
|
* we must use setstate to replace the iterable. One day we
|
|
* will fix pickling of functions
|
|
*/
|
|
if (lz->active) {
|
|
return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
|
|
} else {
|
|
return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
|
|
}
|
|
} else {
|
|
return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
|
|
}
|
|
return NULL;
|
|
}
|
|
|
|
static PyObject *
|
|
chain_setstate(chainobject *lz, PyObject *state)
|
|
{
|
|
PyObject *source, *active=NULL;
|
|
|
|
if (!PyTuple_Check(state)) {
|
|
PyErr_SetString(PyExc_TypeError, "state is not a tuple");
|
|
return NULL;
|
|
}
|
|
if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
|
|
return NULL;
|
|
}
|
|
if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
|
|
PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
|
|
return NULL;
|
|
}
|
|
|
|
Py_INCREF(source);
|
|
Py_XSETREF(lz->source, source);
|
|
Py_XINCREF(active);
|
|
Py_XSETREF(lz->active, active);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
PyDoc_STRVAR(chain_doc,
|
|
"chain(*iterables) --> chain object\n\
|
|
\n\
|
|
Return a chain object whose .__next__() method returns elements from the\n\
|
|
first iterable until it is exhausted, then elements from the next\n\
|
|
iterable, until all of the iterables are exhausted.");
|
|
|
|
static PyMethodDef chain_methods[] = {
|
|
ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
|
|
{"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)chain_setstate, METH_O,
|
|
setstate_doc},
|
|
{"__class_getitem__", Py_GenericAlias,
|
|
METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot chain_slots[] = {
|
|
{Py_tp_dealloc, chain_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)chain_doc},
|
|
{Py_tp_traverse, chain_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, chain_next},
|
|
{Py_tp_methods, chain_methods},
|
|
{Py_tp_new, chain_new},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec chain_spec = {
|
|
.name = "itertools.chain",
|
|
.basicsize = sizeof(chainobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = chain_slots,
|
|
};
|
|
|
|
|
|
/* product object ************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *pools; /* tuple of pool tuples */
|
|
Py_ssize_t *indices; /* one index per pool */
|
|
PyObject *result; /* most recently returned result tuple */
|
|
int stopped; /* set to 1 when the iterator is exhausted */
|
|
} productobject;
|
|
|
|
static PyObject *
|
|
product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
|
{
|
|
productobject *lz;
|
|
Py_ssize_t nargs, npools, repeat=1;
|
|
PyObject *pools = NULL;
|
|
Py_ssize_t *indices = NULL;
|
|
Py_ssize_t i;
|
|
|
|
if (kwds != 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_CheckExact(args));
|
|
if (repeat == 0) {
|
|
nargs = 0;
|
|
} else {
|
|
nargs = PyTuple_GET_SIZE(args);
|
|
if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
|
|
PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
|
|
return NULL;
|
|
}
|
|
}
|
|
npools = nargs * repeat;
|
|
|
|
indices = PyMem_New(Py_ssize_t, npools);
|
|
if (indices == NULL) {
|
|
PyErr_NoMemory();
|
|
goto error;
|
|
}
|
|
|
|
pools = PyTuple_New(npools);
|
|
if (pools == NULL)
|
|
goto error;
|
|
|
|
for (i=0; i < nargs ; ++i) {
|
|
PyObject *item = PyTuple_GET_ITEM(args, i);
|
|
PyObject *pool = PySequence_Tuple(item);
|
|
if (pool == NULL)
|
|
goto error;
|
|
PyTuple_SET_ITEM(pools, i, pool);
|
|
indices[i] = 0;
|
|
}
|
|
for ( ; i < npools; ++i) {
|
|
PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
|
|
Py_INCREF(pool);
|
|
PyTuple_SET_ITEM(pools, i, pool);
|
|
indices[i] = 0;
|
|
}
|
|
|
|
/* create productobject structure */
|
|
lz = (productobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL)
|
|
goto error;
|
|
|
|
lz->pools = pools;
|
|
lz->indices = indices;
|
|
lz->result = NULL;
|
|
lz->stopped = 0;
|
|
|
|
return (PyObject *)lz;
|
|
|
|
error:
|
|
if (indices != NULL)
|
|
PyMem_Free(indices);
|
|
Py_XDECREF(pools);
|
|
return NULL;
|
|
}
|
|
|
|
static void
|
|
product_dealloc(productobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->pools);
|
|
Py_XDECREF(lz->result);
|
|
if (lz->indices != NULL)
|
|
PyMem_Free(lz->indices);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static PyObject *
|
|
product_sizeof(productobject *lz, void *unused)
|
|
{
|
|
size_t res = _PyObject_SIZE(Py_TYPE(lz));
|
|
res += (size_t)PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
|
|
return PyLong_FromSize_t(res);
|
|
}
|
|
|
|
PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
|
|
|
|
static int
|
|
product_traverse(productobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->pools);
|
|
Py_VISIT(lz->result);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
product_next(productobject *lz)
|
|
{
|
|
PyObject *pool;
|
|
PyObject *elem;
|
|
PyObject *oldelem;
|
|
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) {
|
|
/* On the first pass, return an initial tuple filled with the
|
|
first element from each pool. */
|
|
result = PyTuple_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);
|
|
PyTuple_SET_ITEM(result, i, elem);
|
|
}
|
|
} else {
|
|
Py_ssize_t *indices = lz->indices;
|
|
|
|
/* Copy the previous result tuple or re-use it if available */
|
|
if (Py_REFCNT(result) > 1) {
|
|
PyObject *old_result = result;
|
|
result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
|
|
if (result == NULL)
|
|
goto empty;
|
|
lz->result = result;
|
|
Py_DECREF(old_result);
|
|
}
|
|
// bpo-42536: The GC may have untracked this result tuple. Since we're
|
|
// recycling it, make sure it's tracked again:
|
|
else if (!_PyObject_GC_IS_TRACKED(result)) {
|
|
_PyObject_GC_TRACK(result);
|
|
}
|
|
/* Now, we've got the only copy so we can update it in-place */
|
|
assert (npools==0 || 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--) {
|
|
pool = PyTuple_GET_ITEM(pools, i);
|
|
indices[i]++;
|
|
if (indices[i] == PyTuple_GET_SIZE(pool)) {
|
|
/* Roll-over and advance to next pool */
|
|
indices[i] = 0;
|
|
elem = PyTuple_GET_ITEM(pool, 0);
|
|
Py_INCREF(elem);
|
|
oldelem = PyTuple_GET_ITEM(result, i);
|
|
PyTuple_SET_ITEM(result, i, elem);
|
|
Py_DECREF(oldelem);
|
|
} else {
|
|
/* No rollover. Just increment and stop here. */
|
|
elem = PyTuple_GET_ITEM(pool, indices[i]);
|
|
Py_INCREF(elem);
|
|
oldelem = PyTuple_GET_ITEM(result, i);
|
|
PyTuple_SET_ITEM(result, i, elem);
|
|
Py_DECREF(oldelem);
|
|
break;
|
|
}
|
|
}
|
|
|
|
/* If i is negative, then the indices have all rolled-over
|
|
and we're done. */
|
|
if (i < 0)
|
|
goto empty;
|
|
}
|
|
|
|
return Py_NewRef(result);
|
|
|
|
empty:
|
|
lz->stopped = 1;
|
|
return NULL;
|
|
}
|
|
|
|
static PyObject *
|
|
product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
if (lz->stopped) {
|
|
return Py_BuildValue("O(())", Py_TYPE(lz));
|
|
} else if (lz->result == NULL) {
|
|
return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
|
|
} else {
|
|
PyObject *indices;
|
|
Py_ssize_t n, i;
|
|
|
|
/* we must pickle the indices use them for setstate, and
|
|
* additionally indicate that the iterator has started
|
|
*/
|
|
n = PyTuple_GET_SIZE(lz->pools);
|
|
indices = PyTuple_New(n);
|
|
if (indices == NULL)
|
|
return NULL;
|
|
for (i=0; i<n; i++){
|
|
PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
|
|
if (!index) {
|
|
Py_DECREF(indices);
|
|
return NULL;
|
|
}
|
|
PyTuple_SET_ITEM(indices, i, index);
|
|
}
|
|
return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
|
|
}
|
|
}
|
|
|
|
static PyObject *
|
|
product_setstate(productobject *lz, PyObject *state)
|
|
{
|
|
PyObject *result;
|
|
Py_ssize_t n, i;
|
|
|
|
n = PyTuple_GET_SIZE(lz->pools);
|
|
if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
|
|
PyErr_SetString(PyExc_ValueError, "invalid arguments");
|
|
return NULL;
|
|
}
|
|
for (i=0; i<n; i++)
|
|
{
|
|
PyObject* indexObject = PyTuple_GET_ITEM(state, i);
|
|
Py_ssize_t index = PyLong_AsSsize_t(indexObject);
|
|
PyObject* pool;
|
|
Py_ssize_t poolsize;
|
|
if (index < 0 && PyErr_Occurred())
|
|
return NULL; /* not an integer */
|
|
pool = PyTuple_GET_ITEM(lz->pools, i);
|
|
poolsize = PyTuple_GET_SIZE(pool);
|
|
if (poolsize == 0) {
|
|
lz->stopped = 1;
|
|
Py_RETURN_NONE;
|
|
}
|
|
/* clamp the index */
|
|
if (index < 0)
|
|
index = 0;
|
|
else if (index > poolsize-1)
|
|
index = poolsize-1;
|
|
lz->indices[i] = index;
|
|
}
|
|
|
|
result = PyTuple_New(n);
|
|
if (!result)
|
|
return NULL;
|
|
for (i=0; i<n; i++) {
|
|
PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
|
|
PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
|
|
Py_INCREF(element);
|
|
PyTuple_SET_ITEM(result, i, element);
|
|
}
|
|
Py_XSETREF(lz->result, result);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef product_methods[] = {
|
|
{"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)product_setstate, METH_O,
|
|
setstate_doc},
|
|
{"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
|
|
sizeof_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
PyDoc_STRVAR(product_doc,
|
|
"product(*iterables, repeat=1) --> product object\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\
|
|
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\
|
|
To compute the product of an iterable with itself, specify the number\n\
|
|
of repetitions with the optional repeat keyword argument. For example,\n\
|
|
product(A, repeat=4) means the same as product(A, A, A, A).\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 PyType_Slot product_slots[] = {
|
|
{Py_tp_dealloc, product_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)product_doc},
|
|
{Py_tp_traverse, product_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, product_next},
|
|
{Py_tp_methods, product_methods},
|
|
{Py_tp_new, product_new},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec product_spec = {
|
|
.name = "itertools.product",
|
|
.basicsize = sizeof(productobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = product_slots,
|
|
};
|
|
|
|
|
|
/* combinations object *******************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *pool; /* input converted to a tuple */
|
|
Py_ssize_t *indices; /* one index per result element */
|
|
PyObject *result; /* most recently returned result tuple */
|
|
Py_ssize_t r; /* size of result tuple */
|
|
int stopped; /* set to 1 when the iterator is exhausted */
|
|
} combinationsobject;
|
|
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.combinations.__new__
|
|
iterable: object
|
|
r: Py_ssize_t
|
|
Return successive r-length combinations of elements in the iterable.
|
|
|
|
combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
|
|
Py_ssize_t r)
|
|
/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
|
|
{
|
|
combinationsobject *co;
|
|
Py_ssize_t n;
|
|
PyObject *pool = NULL;
|
|
Py_ssize_t *indices = NULL;
|
|
Py_ssize_t i;
|
|
|
|
pool = PySequence_Tuple(iterable);
|
|
if (pool == NULL)
|
|
goto error;
|
|
n = PyTuple_GET_SIZE(pool);
|
|
if (r < 0) {
|
|
PyErr_SetString(PyExc_ValueError, "r must be non-negative");
|
|
goto error;
|
|
}
|
|
|
|
indices = PyMem_New(Py_ssize_t, r);
|
|
if (indices == NULL) {
|
|
PyErr_NoMemory();
|
|
goto error;
|
|
}
|
|
|
|
for (i=0 ; i<r ; i++)
|
|
indices[i] = i;
|
|
|
|
/* create combinationsobject structure */
|
|
co = (combinationsobject *)type->tp_alloc(type, 0);
|
|
if (co == NULL)
|
|
goto error;
|
|
|
|
co->pool = pool;
|
|
co->indices = indices;
|
|
co->result = NULL;
|
|
co->r = r;
|
|
co->stopped = r > n ? 1 : 0;
|
|
|
|
return (PyObject *)co;
|
|
|
|
error:
|
|
if (indices != NULL)
|
|
PyMem_Free(indices);
|
|
Py_XDECREF(pool);
|
|
return NULL;
|
|
}
|
|
|
|
static void
|
|
combinations_dealloc(combinationsobject *co)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(co);
|
|
PyObject_GC_UnTrack(co);
|
|
Py_XDECREF(co->pool);
|
|
Py_XDECREF(co->result);
|
|
if (co->indices != NULL)
|
|
PyMem_Free(co->indices);
|
|
tp->tp_free(co);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static PyObject *
|
|
combinations_sizeof(combinationsobject *co, void *unused)
|
|
{
|
|
size_t res = _PyObject_SIZE(Py_TYPE(co));
|
|
res += (size_t)co->r * sizeof(Py_ssize_t);
|
|
return PyLong_FromSize_t(res);
|
|
}
|
|
|
|
static int
|
|
combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(co));
|
|
Py_VISIT(co->pool);
|
|
Py_VISIT(co->result);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
combinations_next(combinationsobject *co)
|
|
{
|
|
PyObject *elem;
|
|
PyObject *oldelem;
|
|
PyObject *pool = co->pool;
|
|
Py_ssize_t *indices = co->indices;
|
|
PyObject *result = co->result;
|
|
Py_ssize_t n = PyTuple_GET_SIZE(pool);
|
|
Py_ssize_t r = co->r;
|
|
Py_ssize_t i, j, index;
|
|
|
|
if (co->stopped)
|
|
return NULL;
|
|
|
|
if (result == NULL) {
|
|
/* On the first pass, initialize result tuple using the indices */
|
|
result = PyTuple_New(r);
|
|
if (result == NULL)
|
|
goto empty;
|
|
co->result = result;
|
|
for (i=0; i<r ; i++) {
|
|
index = indices[i];
|
|
elem = PyTuple_GET_ITEM(pool, index);
|
|
Py_INCREF(elem);
|
|
PyTuple_SET_ITEM(result, i, elem);
|
|
}
|
|
} else {
|
|
/* Copy the previous result tuple or re-use it if available */
|
|
if (Py_REFCNT(result) > 1) {
|
|
PyObject *old_result = result;
|
|
result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
|
|
if (result == NULL)
|
|
goto empty;
|
|
co->result = result;
|
|
Py_DECREF(old_result);
|
|
}
|
|
// bpo-42536: The GC may have untracked this result tuple. Since we're
|
|
// recycling it, make sure it's tracked again:
|
|
else if (!_PyObject_GC_IS_TRACKED(result)) {
|
|
_PyObject_GC_TRACK(result);
|
|
}
|
|
/* Now, we've got the only copy so we can update it in-place
|
|
* CPython's empty tuple is a singleton and cached in
|
|
* PyTuple's freelist.
|
|
*/
|
|
assert(r == 0 || Py_REFCNT(result) == 1);
|
|
|
|
/* Scan indices right-to-left until finding one that is not
|
|
at its maximum (i + n - r). */
|
|
for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
|
|
;
|
|
|
|
/* If i is negative, then the indices are all at
|
|
their maximum value and we're done. */
|
|
if (i < 0)
|
|
goto empty;
|
|
|
|
/* Increment the current index which we know is not at its
|
|
maximum. Then move back to the right setting each index
|
|
to its lowest possible value (one higher than the index
|
|
to its left -- this maintains the sort order invariant). */
|
|
indices[i]++;
|
|
for (j=i+1 ; j<r ; j++)
|
|
indices[j] = indices[j-1] + 1;
|
|
|
|
/* Update the result tuple for the new indices
|
|
starting with i, the leftmost index that changed */
|
|
for ( ; i<r ; i++) {
|
|
index = indices[i];
|
|
elem = PyTuple_GET_ITEM(pool, index);
|
|
Py_INCREF(elem);
|
|
oldelem = PyTuple_GET_ITEM(result, i);
|
|
PyTuple_SET_ITEM(result, i, elem);
|
|
Py_DECREF(oldelem);
|
|
}
|
|
}
|
|
|
|
return Py_NewRef(result);
|
|
|
|
empty:
|
|
co->stopped = 1;
|
|
return NULL;
|
|
}
|
|
|
|
static PyObject *
|
|
combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
if (lz->result == NULL) {
|
|
return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
|
|
} else if (lz->stopped) {
|
|
return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
|
|
} else {
|
|
PyObject *indices;
|
|
Py_ssize_t i;
|
|
|
|
/* we must pickle the indices and use them for setstate */
|
|
indices = PyTuple_New(lz->r);
|
|
if (!indices)
|
|
return NULL;
|
|
for (i=0; i<lz->r; i++)
|
|
{
|
|
PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
|
|
if (!index) {
|
|
Py_DECREF(indices);
|
|
return NULL;
|
|
}
|
|
PyTuple_SET_ITEM(indices, i, index);
|
|
}
|
|
|
|
return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
|
|
}
|
|
}
|
|
|
|
static PyObject *
|
|
combinations_setstate(combinationsobject *lz, PyObject *state)
|
|
{
|
|
PyObject *result;
|
|
Py_ssize_t i;
|
|
Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
|
|
|
|
if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
|
|
PyErr_SetString(PyExc_ValueError, "invalid arguments");
|
|
return NULL;
|
|
}
|
|
|
|
for (i=0; i<lz->r; i++) {
|
|
Py_ssize_t max;
|
|
PyObject* indexObject = PyTuple_GET_ITEM(state, i);
|
|
Py_ssize_t index = PyLong_AsSsize_t(indexObject);
|
|
|
|
if (index == -1 && PyErr_Occurred())
|
|
return NULL; /* not an integer */
|
|
max = i + n - lz->r;
|
|
/* clamp the index (beware of negative max) */
|
|
if (index > max)
|
|
index = max;
|
|
if (index < 0)
|
|
index = 0;
|
|
lz->indices[i] = index;
|
|
}
|
|
|
|
result = PyTuple_New(lz->r);
|
|
if (result == NULL)
|
|
return NULL;
|
|
for (i=0; i<lz->r; i++) {
|
|
PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
|
|
Py_INCREF(element);
|
|
PyTuple_SET_ITEM(result, i, element);
|
|
}
|
|
|
|
Py_XSETREF(lz->result, result);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef combinations_methods[] = {
|
|
{"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)combinations_setstate, METH_O,
|
|
setstate_doc},
|
|
{"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
|
|
sizeof_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot combinations_slots[] = {
|
|
{Py_tp_dealloc, combinations_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_combinations__doc__},
|
|
{Py_tp_traverse, combinations_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, combinations_next},
|
|
{Py_tp_methods, combinations_methods},
|
|
{Py_tp_new, itertools_combinations},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec combinations_spec = {
|
|
.name = "itertools.combinations",
|
|
.basicsize = sizeof(combinationsobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = combinations_slots,
|
|
};
|
|
|
|
|
|
/* combinations with replacement object **************************************/
|
|
|
|
/* Equivalent to:
|
|
|
|
def combinations_with_replacement(iterable, r):
|
|
"combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
|
|
# number items returned: (n+r-1)! / r! / (n-1)!
|
|
pool = tuple(iterable)
|
|
n = len(pool)
|
|
indices = [0] * r
|
|
yield tuple(pool[i] for i in indices)
|
|
while 1:
|
|
for i in reversed(range(r)):
|
|
if indices[i] != n - 1:
|
|
break
|
|
else:
|
|
return
|
|
indices[i:] = [indices[i] + 1] * (r - i)
|
|
yield tuple(pool[i] for i in indices)
|
|
|
|
def combinations_with_replacement2(iterable, r):
|
|
'Alternate version that filters from product()'
|
|
pool = tuple(iterable)
|
|
n = len(pool)
|
|
for indices in product(range(n), repeat=r):
|
|
if sorted(indices) == list(indices):
|
|
yield tuple(pool[i] for i in indices)
|
|
*/
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *pool; /* input converted to a tuple */
|
|
Py_ssize_t *indices; /* one index per result element */
|
|
PyObject *result; /* most recently returned result tuple */
|
|
Py_ssize_t r; /* size of result tuple */
|
|
int stopped; /* set to 1 when the cwr iterator is exhausted */
|
|
} cwrobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.combinations_with_replacement.__new__
|
|
iterable: object
|
|
r: Py_ssize_t
|
|
Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
|
|
|
|
combinations_with_replacement('ABC', 2) --> ('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_combinations_with_replacement_impl(PyTypeObject *type,
|
|
PyObject *iterable,
|
|
Py_ssize_t r)
|
|
/*[clinic end generated code: output=48b26856d4e659ca input=1dc58e82a0878fdc]*/
|
|
{
|
|
cwrobject *co;
|
|
Py_ssize_t n;
|
|
PyObject *pool = NULL;
|
|
Py_ssize_t *indices = NULL;
|
|
Py_ssize_t i;
|
|
|
|
pool = PySequence_Tuple(iterable);
|
|
if (pool == NULL)
|
|
goto error;
|
|
n = PyTuple_GET_SIZE(pool);
|
|
if (r < 0) {
|
|
PyErr_SetString(PyExc_ValueError, "r must be non-negative");
|
|
goto error;
|
|
}
|
|
|
|
indices = PyMem_New(Py_ssize_t, r);
|
|
if (indices == NULL) {
|
|
PyErr_NoMemory();
|
|
goto error;
|
|
}
|
|
|
|
for (i=0 ; i<r ; i++)
|
|
indices[i] = 0;
|
|
|
|
/* create cwrobject structure */
|
|
co = (cwrobject *)type->tp_alloc(type, 0);
|
|
if (co == NULL)
|
|
goto error;
|
|
|
|
co->pool = pool;
|
|
co->indices = indices;
|
|
co->result = NULL;
|
|
co->r = r;
|
|
co->stopped = !n && r;
|
|
|
|
return (PyObject *)co;
|
|
|
|
error:
|
|
if (indices != NULL)
|
|
PyMem_Free(indices);
|
|
Py_XDECREF(pool);
|
|
return NULL;
|
|
}
|
|
|
|
static void
|
|
cwr_dealloc(cwrobject *co)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(co);
|
|
PyObject_GC_UnTrack(co);
|
|
Py_XDECREF(co->pool);
|
|
Py_XDECREF(co->result);
|
|
if (co->indices != NULL)
|
|
PyMem_Free(co->indices);
|
|
tp->tp_free(co);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static PyObject *
|
|
cwr_sizeof(cwrobject *co, void *unused)
|
|
{
|
|
size_t res = _PyObject_SIZE(Py_TYPE(co));
|
|
res += (size_t)co->r * sizeof(Py_ssize_t);
|
|
return PyLong_FromSize_t(res);
|
|
}
|
|
|
|
static int
|
|
cwr_traverse(cwrobject *co, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(co));
|
|
Py_VISIT(co->pool);
|
|
Py_VISIT(co->result);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
cwr_next(cwrobject *co)
|
|
{
|
|
PyObject *elem;
|
|
PyObject *oldelem;
|
|
PyObject *pool = co->pool;
|
|
Py_ssize_t *indices = co->indices;
|
|
PyObject *result = co->result;
|
|
Py_ssize_t n = PyTuple_GET_SIZE(pool);
|
|
Py_ssize_t r = co->r;
|
|
Py_ssize_t i, index;
|
|
|
|
if (co->stopped)
|
|
return NULL;
|
|
|
|
if (result == NULL) {
|
|
/* On the first pass, initialize result tuple with pool[0] */
|
|
result = PyTuple_New(r);
|
|
if (result == NULL)
|
|
goto empty;
|
|
co->result = result;
|
|
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 */
|
|
if (Py_REFCNT(result) > 1) {
|
|
PyObject *old_result = result;
|
|
result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
|
|
if (result == NULL)
|
|
goto empty;
|
|
co->result = result;
|
|
Py_DECREF(old_result);
|
|
}
|
|
// bpo-42536: The GC may have untracked this result tuple. Since we're
|
|
// recycling it, make sure it's tracked again:
|
|
else if (!_PyObject_GC_IS_TRACKED(result)) {
|
|
_PyObject_GC_TRACK(result);
|
|
}
|
|
/* Now, we've got the only copy so we can update it in-place CPython's
|
|
empty tuple is a singleton and cached in PyTuple's freelist. */
|
|
assert(r == 0 || Py_REFCNT(result) == 1);
|
|
|
|
/* Scan indices right-to-left until finding one that is not
|
|
* at its maximum (n-1). */
|
|
for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
|
|
;
|
|
|
|
/* If i is negative, then the indices are all at
|
|
their maximum value and we're done. */
|
|
if (i < 0)
|
|
goto empty;
|
|
|
|
/* Increment the current index which we know is not at its
|
|
maximum. Then set all to the right to the same value. */
|
|
index = indices[i] + 1;
|
|
assert(index < n);
|
|
elem = PyTuple_GET_ITEM(pool, index);
|
|
for ( ; i<r ; i++) {
|
|
indices[i] = index;
|
|
Py_INCREF(elem);
|
|
oldelem = PyTuple_GET_ITEM(result, i);
|
|
PyTuple_SET_ITEM(result, i, elem);
|
|
Py_DECREF(oldelem);
|
|
}
|
|
}
|
|
|
|
return Py_NewRef(result);
|
|
|
|
empty:
|
|
co->stopped = 1;
|
|
return NULL;
|
|
}
|
|
|
|
static PyObject *
|
|
cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
if (lz->result == NULL) {
|
|
return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
|
|
} else if (lz->stopped) {
|
|
return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
|
|
} else {
|
|
PyObject *indices;
|
|
Py_ssize_t i;
|
|
|
|
/* we must pickle the indices and use them for setstate */
|
|
indices = PyTuple_New(lz->r);
|
|
if (!indices)
|
|
return NULL;
|
|
for (i=0; i<lz->r; i++) {
|
|
PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
|
|
if (!index) {
|
|
Py_DECREF(indices);
|
|
return NULL;
|
|
}
|
|
PyTuple_SET_ITEM(indices, i, index);
|
|
}
|
|
|
|
return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
|
|
}
|
|
}
|
|
|
|
static PyObject *
|
|
cwr_setstate(cwrobject *lz, PyObject *state)
|
|
{
|
|
PyObject *result;
|
|
Py_ssize_t n, i;
|
|
|
|
if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
|
|
{
|
|
PyErr_SetString(PyExc_ValueError, "invalid arguments");
|
|
return NULL;
|
|
}
|
|
|
|
n = PyTuple_GET_SIZE(lz->pool);
|
|
for (i=0; i<lz->r; i++) {
|
|
PyObject* indexObject = PyTuple_GET_ITEM(state, i);
|
|
Py_ssize_t index = PyLong_AsSsize_t(indexObject);
|
|
|
|
if (index < 0 && PyErr_Occurred())
|
|
return NULL; /* not an integer */
|
|
/* clamp the index */
|
|
if (index < 0)
|
|
index = 0;
|
|
else if (index > n-1)
|
|
index = n-1;
|
|
lz->indices[i] = index;
|
|
}
|
|
result = PyTuple_New(lz->r);
|
|
if (result == NULL)
|
|
return NULL;
|
|
for (i=0; i<lz->r; i++) {
|
|
PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
|
|
Py_INCREF(element);
|
|
PyTuple_SET_ITEM(result, i, element);
|
|
}
|
|
Py_XSETREF(lz->result, result);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef cwr_methods[] = {
|
|
{"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)cwr_setstate, METH_O,
|
|
setstate_doc},
|
|
{"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
|
|
sizeof_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot cwr_slots[] = {
|
|
{Py_tp_dealloc, cwr_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_combinations_with_replacement__doc__},
|
|
{Py_tp_traverse, cwr_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, cwr_next},
|
|
{Py_tp_methods, cwr_methods},
|
|
{Py_tp_new, itertools_combinations_with_replacement},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec cwr_spec = {
|
|
.name = "itertools.combinations_with_replacement",
|
|
.basicsize = sizeof(cwrobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = cwr_slots,
|
|
};
|
|
|
|
|
|
/* permutations object ********************************************************
|
|
|
|
def permutations(iterable, r=None):
|
|
# permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
|
|
# permutations(range(3)) --> 012 021 102 120 201 210
|
|
pool = tuple(iterable)
|
|
n = len(pool)
|
|
r = n if r is None else r
|
|
if r > n:
|
|
return
|
|
indices = list(range(n))
|
|
cycles = list(range(n, n-r, -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
|
|
else:
|
|
j = cycles[i]
|
|
indices[i], indices[-j] = indices[-j], indices[i]
|
|
yield tuple(pool[i] for i in indices[:r])
|
|
break
|
|
else:
|
|
return
|
|
*/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *pool; /* input converted to a tuple */
|
|
Py_ssize_t *indices; /* one index per element in the pool */
|
|
Py_ssize_t *cycles; /* one rollover counter per element in the result */
|
|
PyObject *result; /* most recently returned result tuple */
|
|
Py_ssize_t r; /* size of result tuple */
|
|
int stopped; /* set to 1 when the iterator is exhausted */
|
|
} permutationsobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.permutations.__new__
|
|
iterable: object
|
|
r as robj: object = None
|
|
Return successive r-length permutations of elements in the iterable.
|
|
|
|
permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
|
|
PyObject *robj)
|
|
/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
|
|
{
|
|
permutationsobject *po;
|
|
Py_ssize_t n;
|
|
Py_ssize_t r;
|
|
PyObject *pool = NULL;
|
|
Py_ssize_t *indices = NULL;
|
|
Py_ssize_t *cycles = NULL;
|
|
Py_ssize_t i;
|
|
|
|
pool = PySequence_Tuple(iterable);
|
|
if (pool == NULL)
|
|
goto error;
|
|
n = PyTuple_GET_SIZE(pool);
|
|
|
|
r = n;
|
|
if (robj != Py_None) {
|
|
if (!PyLong_Check(robj)) {
|
|
PyErr_SetString(PyExc_TypeError, "Expected int as r");
|
|
goto error;
|
|
}
|
|
r = PyLong_AsSsize_t(robj);
|
|
if (r == -1 && PyErr_Occurred())
|
|
goto error;
|
|
}
|
|
if (r < 0) {
|
|
PyErr_SetString(PyExc_ValueError, "r must be non-negative");
|
|
goto error;
|
|
}
|
|
|
|
indices = PyMem_New(Py_ssize_t, n);
|
|
cycles = PyMem_New(Py_ssize_t, r);
|
|
if (indices == NULL || cycles == NULL) {
|
|
PyErr_NoMemory();
|
|
goto error;
|
|
}
|
|
|
|
for (i=0 ; i<n ; i++)
|
|
indices[i] = i;
|
|
for (i=0 ; i<r ; i++)
|
|
cycles[i] = n - i;
|
|
|
|
/* create permutationsobject structure */
|
|
po = (permutationsobject *)type->tp_alloc(type, 0);
|
|
if (po == NULL)
|
|
goto error;
|
|
|
|
po->pool = pool;
|
|
po->indices = indices;
|
|
po->cycles = cycles;
|
|
po->result = NULL;
|
|
po->r = r;
|
|
po->stopped = r > n ? 1 : 0;
|
|
|
|
return (PyObject *)po;
|
|
|
|
error:
|
|
if (indices != NULL)
|
|
PyMem_Free(indices);
|
|
if (cycles != NULL)
|
|
PyMem_Free(cycles);
|
|
Py_XDECREF(pool);
|
|
return NULL;
|
|
}
|
|
|
|
static void
|
|
permutations_dealloc(permutationsobject *po)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(po);
|
|
PyObject_GC_UnTrack(po);
|
|
Py_XDECREF(po->pool);
|
|
Py_XDECREF(po->result);
|
|
PyMem_Free(po->indices);
|
|
PyMem_Free(po->cycles);
|
|
tp->tp_free(po);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static PyObject *
|
|
permutations_sizeof(permutationsobject *po, void *unused)
|
|
{
|
|
size_t res = _PyObject_SIZE(Py_TYPE(po));
|
|
res += (size_t)PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
|
|
res += (size_t)po->r * sizeof(Py_ssize_t);
|
|
return PyLong_FromSize_t(res);
|
|
}
|
|
|
|
static int
|
|
permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(po));
|
|
Py_VISIT(po->pool);
|
|
Py_VISIT(po->result);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
permutations_next(permutationsobject *po)
|
|
{
|
|
PyObject *elem;
|
|
PyObject *oldelem;
|
|
PyObject *pool = po->pool;
|
|
Py_ssize_t *indices = po->indices;
|
|
Py_ssize_t *cycles = po->cycles;
|
|
PyObject *result = po->result;
|
|
Py_ssize_t n = PyTuple_GET_SIZE(pool);
|
|
Py_ssize_t r = po->r;
|
|
Py_ssize_t i, j, k, index;
|
|
|
|
if (po->stopped)
|
|
return NULL;
|
|
|
|
if (result == NULL) {
|
|
/* On the first pass, initialize result tuple using the indices */
|
|
result = PyTuple_New(r);
|
|
if (result == NULL)
|
|
goto empty;
|
|
po->result = result;
|
|
for (i=0; i<r ; i++) {
|
|
index = indices[i];
|
|
elem = PyTuple_GET_ITEM(pool, index);
|
|
Py_INCREF(elem);
|
|
PyTuple_SET_ITEM(result, i, elem);
|
|
}
|
|
} else {
|
|
if (n == 0)
|
|
goto empty;
|
|
|
|
/* Copy the previous result tuple or re-use it if available */
|
|
if (Py_REFCNT(result) > 1) {
|
|
PyObject *old_result = result;
|
|
result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
|
|
if (result == NULL)
|
|
goto empty;
|
|
po->result = result;
|
|
Py_DECREF(old_result);
|
|
}
|
|
// bpo-42536: The GC may have untracked this result tuple. Since we're
|
|
// recycling it, make sure it's tracked again:
|
|
else if (!_PyObject_GC_IS_TRACKED(result)) {
|
|
_PyObject_GC_TRACK(result);
|
|
}
|
|
/* Now, we've got the only copy so we can update it in-place */
|
|
assert(r == 0 || Py_REFCNT(result) == 1);
|
|
|
|
/* Decrement rightmost cycle, moving leftward upon zero rollover */
|
|
for (i=r-1 ; i>=0 ; i--) {
|
|
cycles[i] -= 1;
|
|
if (cycles[i] == 0) {
|
|
/* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
|
|
index = indices[i];
|
|
for (j=i ; j<n-1 ; j++)
|
|
indices[j] = indices[j+1];
|
|
indices[n-1] = index;
|
|
cycles[i] = n - i;
|
|
} else {
|
|
j = cycles[i];
|
|
index = indices[i];
|
|
indices[i] = indices[n-j];
|
|
indices[n-j] = index;
|
|
|
|
for (k=i; k<r ; k++) {
|
|
/* start with i, the leftmost element that changed */
|
|
/* yield tuple(pool[k] for k in indices[:r]) */
|
|
index = indices[k];
|
|
elem = PyTuple_GET_ITEM(pool, index);
|
|
Py_INCREF(elem);
|
|
oldelem = PyTuple_GET_ITEM(result, k);
|
|
PyTuple_SET_ITEM(result, k, elem);
|
|
Py_DECREF(oldelem);
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
/* If i is negative, then the cycles have all
|
|
rolled-over and we're done. */
|
|
if (i < 0)
|
|
goto empty;
|
|
}
|
|
return Py_NewRef(result);
|
|
|
|
empty:
|
|
po->stopped = 1;
|
|
return NULL;
|
|
}
|
|
|
|
static PyObject *
|
|
permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
if (po->result == NULL) {
|
|
return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
|
|
} else if (po->stopped) {
|
|
return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
|
|
} else {
|
|
PyObject *indices=NULL, *cycles=NULL;
|
|
Py_ssize_t n, i;
|
|
|
|
/* we must pickle the indices and cycles and use them for setstate */
|
|
n = PyTuple_GET_SIZE(po->pool);
|
|
indices = PyTuple_New(n);
|
|
if (indices == NULL)
|
|
goto err;
|
|
for (i=0; i<n; i++) {
|
|
PyObject* index = PyLong_FromSsize_t(po->indices[i]);
|
|
if (!index)
|
|
goto err;
|
|
PyTuple_SET_ITEM(indices, i, index);
|
|
}
|
|
|
|
cycles = PyTuple_New(po->r);
|
|
if (cycles == NULL)
|
|
goto err;
|
|
for (i=0 ; i<po->r ; i++) {
|
|
PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
|
|
if (!index)
|
|
goto err;
|
|
PyTuple_SET_ITEM(cycles, i, index);
|
|
}
|
|
return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
|
|
po->pool, po->r,
|
|
indices, cycles);
|
|
err:
|
|
Py_XDECREF(indices);
|
|
Py_XDECREF(cycles);
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
static PyObject *
|
|
permutations_setstate(permutationsobject *po, PyObject *state)
|
|
{
|
|
PyObject *indices, *cycles, *result;
|
|
Py_ssize_t n, i;
|
|
|
|
if (!PyTuple_Check(state)) {
|
|
PyErr_SetString(PyExc_TypeError, "state is not a tuple");
|
|
return NULL;
|
|
}
|
|
if (!PyArg_ParseTuple(state, "O!O!",
|
|
&PyTuple_Type, &indices,
|
|
&PyTuple_Type, &cycles)) {
|
|
return NULL;
|
|
}
|
|
|
|
n = PyTuple_GET_SIZE(po->pool);
|
|
if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
|
|
PyErr_SetString(PyExc_ValueError, "invalid arguments");
|
|
return NULL;
|
|
}
|
|
|
|
for (i=0; i<n; i++) {
|
|
PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
|
|
Py_ssize_t index = PyLong_AsSsize_t(indexObject);
|
|
if (index < 0 && PyErr_Occurred())
|
|
return NULL; /* not an integer */
|
|
/* clamp the index */
|
|
if (index < 0)
|
|
index = 0;
|
|
else if (index > n-1)
|
|
index = n-1;
|
|
po->indices[i] = index;
|
|
}
|
|
|
|
for (i=0; i<po->r; i++) {
|
|
PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
|
|
Py_ssize_t index = PyLong_AsSsize_t(indexObject);
|
|
if (index < 0 && PyErr_Occurred())
|
|
return NULL; /* not an integer */
|
|
if (index < 1)
|
|
index = 1;
|
|
else if (index > n-i)
|
|
index = n-i;
|
|
po->cycles[i] = index;
|
|
}
|
|
result = PyTuple_New(po->r);
|
|
if (result == NULL)
|
|
return NULL;
|
|
for (i=0; i<po->r; i++) {
|
|
PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
|
|
Py_INCREF(element);
|
|
PyTuple_SET_ITEM(result, i, element);
|
|
}
|
|
Py_XSETREF(po->result, result);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef permuations_methods[] = {
|
|
{"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)permutations_setstate, METH_O,
|
|
setstate_doc},
|
|
{"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
|
|
sizeof_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot permutations_slots[] = {
|
|
{Py_tp_dealloc, permutations_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_permutations__doc__},
|
|
{Py_tp_traverse, permutations_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, permutations_next},
|
|
{Py_tp_methods, permuations_methods},
|
|
{Py_tp_new, itertools_permutations},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec permutations_spec = {
|
|
.name = "itertools.permutations",
|
|
.basicsize = sizeof(permutationsobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = permutations_slots,
|
|
};
|
|
|
|
|
|
/* accumulate object ********************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *total;
|
|
PyObject *it;
|
|
PyObject *binop;
|
|
PyObject *initial;
|
|
itertools_state *state;
|
|
} accumulateobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.accumulate.__new__
|
|
iterable: object
|
|
func as binop: object = None
|
|
*
|
|
initial: object = None
|
|
Return series of accumulated sums (or other binary function results).
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
|
|
PyObject *binop, PyObject *initial)
|
|
/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
|
|
{
|
|
PyObject *it;
|
|
accumulateobject *lz;
|
|
|
|
/* Get iterator. */
|
|
it = PyObject_GetIter(iterable);
|
|
if (it == NULL)
|
|
return NULL;
|
|
|
|
/* create accumulateobject structure */
|
|
lz = (accumulateobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
|
|
if (binop != Py_None) {
|
|
lz->binop = Py_XNewRef(binop);
|
|
}
|
|
lz->total = NULL;
|
|
lz->it = it;
|
|
lz->initial = Py_XNewRef(initial);
|
|
lz->state = find_state_by_type(type);
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
accumulate_dealloc(accumulateobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->binop);
|
|
Py_XDECREF(lz->total);
|
|
Py_XDECREF(lz->it);
|
|
Py_XDECREF(lz->initial);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->binop);
|
|
Py_VISIT(lz->it);
|
|
Py_VISIT(lz->total);
|
|
Py_VISIT(lz->initial);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
accumulate_next(accumulateobject *lz)
|
|
{
|
|
PyObject *val, *newtotal;
|
|
|
|
if (lz->initial != Py_None) {
|
|
lz->total = lz->initial;
|
|
lz->initial = Py_NewRef(Py_None);
|
|
return Py_NewRef(lz->total);
|
|
}
|
|
val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
|
|
if (val == NULL)
|
|
return NULL;
|
|
|
|
if (lz->total == NULL) {
|
|
lz->total = Py_NewRef(val);
|
|
return lz->total;
|
|
}
|
|
|
|
if (lz->binop == NULL)
|
|
newtotal = PyNumber_Add(lz->total, val);
|
|
else
|
|
newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
|
|
Py_DECREF(val);
|
|
if (newtotal == NULL)
|
|
return NULL;
|
|
|
|
Py_INCREF(newtotal);
|
|
Py_SETREF(lz->total, newtotal);
|
|
return newtotal;
|
|
}
|
|
|
|
static PyObject *
|
|
accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
itertools_state *state = lz->state;
|
|
|
|
if (lz->initial != Py_None) {
|
|
PyObject *it;
|
|
|
|
assert(lz->total == NULL);
|
|
it = PyObject_CallFunction((PyObject *)(state->chain_type), "(O)O",
|
|
lz->initial, lz->it);
|
|
if (it == NULL)
|
|
return NULL;
|
|
return Py_BuildValue("O(NO)O", Py_TYPE(lz),
|
|
it, lz->binop?lz->binop:Py_None, Py_None);
|
|
}
|
|
if (lz->total == Py_None) {
|
|
PyObject *it;
|
|
|
|
it = PyObject_CallFunction((PyObject *)(state->chain_type), "(O)O",
|
|
lz->total, lz->it);
|
|
if (it == NULL)
|
|
return NULL;
|
|
it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
|
|
it, lz->binop ? lz->binop : Py_None);
|
|
if (it == NULL)
|
|
return NULL;
|
|
|
|
return Py_BuildValue("O(NiO)", state->islice_type, it, 1, Py_None);
|
|
}
|
|
return Py_BuildValue("O(OO)O", Py_TYPE(lz),
|
|
lz->it, lz->binop?lz->binop:Py_None,
|
|
lz->total?lz->total:Py_None);
|
|
}
|
|
|
|
static PyObject *
|
|
accumulate_setstate(accumulateobject *lz, PyObject *state)
|
|
{
|
|
Py_INCREF(state);
|
|
Py_XSETREF(lz->total, state);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef accumulate_methods[] = {
|
|
{"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
|
|
setstate_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot accumulate_slots[] = {
|
|
{Py_tp_dealloc, accumulate_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_accumulate__doc__},
|
|
{Py_tp_traverse, accumulate_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, accumulate_next},
|
|
{Py_tp_methods, accumulate_methods},
|
|
{Py_tp_new, itertools_accumulate},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec accumulate_spec = {
|
|
.name = "itertools.accumulate",
|
|
.basicsize = sizeof(accumulateobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = accumulate_slots,
|
|
};
|
|
|
|
|
|
/* compress object ************************************************************/
|
|
|
|
/* Equivalent to:
|
|
|
|
def compress(data, selectors):
|
|
"compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
|
|
return (d for d, s in zip(data, selectors) if s)
|
|
*/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *data;
|
|
PyObject *selectors;
|
|
} compressobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.compress.__new__
|
|
data as seq1: object
|
|
selectors as seq2: object
|
|
Return data elements corresponding to true selector elements.
|
|
|
|
Forms a shorter iterator from selected data elements using the selectors to
|
|
choose the data elements.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
|
|
/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
|
|
{
|
|
PyObject *data=NULL, *selectors=NULL;
|
|
compressobject *lz;
|
|
|
|
data = PyObject_GetIter(seq1);
|
|
if (data == NULL)
|
|
goto fail;
|
|
selectors = PyObject_GetIter(seq2);
|
|
if (selectors == NULL)
|
|
goto fail;
|
|
|
|
/* create compressobject structure */
|
|
lz = (compressobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL)
|
|
goto fail;
|
|
lz->data = data;
|
|
lz->selectors = selectors;
|
|
return (PyObject *)lz;
|
|
|
|
fail:
|
|
Py_XDECREF(data);
|
|
Py_XDECREF(selectors);
|
|
return NULL;
|
|
}
|
|
|
|
static void
|
|
compress_dealloc(compressobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->data);
|
|
Py_XDECREF(lz->selectors);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
compress_traverse(compressobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->data);
|
|
Py_VISIT(lz->selectors);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
compress_next(compressobject *lz)
|
|
{
|
|
PyObject *data = lz->data, *selectors = lz->selectors;
|
|
PyObject *datum, *selector;
|
|
PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
|
|
PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
|
|
int ok;
|
|
|
|
while (1) {
|
|
/* Steps: get datum, get selector, evaluate selector.
|
|
Order is important (to match the pure python version
|
|
in terms of which input gets a chance to raise an
|
|
exception first).
|
|
*/
|
|
|
|
datum = datanext(data);
|
|
if (datum == NULL)
|
|
return NULL;
|
|
|
|
selector = selectornext(selectors);
|
|
if (selector == NULL) {
|
|
Py_DECREF(datum);
|
|
return NULL;
|
|
}
|
|
|
|
ok = PyObject_IsTrue(selector);
|
|
Py_DECREF(selector);
|
|
if (ok > 0)
|
|
return datum;
|
|
Py_DECREF(datum);
|
|
if (ok < 0)
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
static PyObject *
|
|
compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
return Py_BuildValue("O(OO)", Py_TYPE(lz),
|
|
lz->data, lz->selectors);
|
|
}
|
|
|
|
static PyMethodDef compress_methods[] = {
|
|
{"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot compress_slots[] = {
|
|
{Py_tp_dealloc, compress_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_compress__doc__},
|
|
{Py_tp_traverse, compress_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, compress_next},
|
|
{Py_tp_methods, compress_methods},
|
|
{Py_tp_new, itertools_compress},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec compress_spec = {
|
|
.name = "itertools.compress",
|
|
.basicsize = sizeof(compressobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = compress_slots,
|
|
};
|
|
|
|
|
|
/* filterfalse object ************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *func;
|
|
PyObject *it;
|
|
} filterfalseobject;
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.filterfalse.__new__
|
|
function as func: object
|
|
iterable as seq: object
|
|
/
|
|
Return those items of iterable for which function(item) is false.
|
|
|
|
If function is None, return the items that are false.
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
|
|
/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
|
|
{
|
|
PyObject *it;
|
|
filterfalseobject *lz;
|
|
|
|
/* Get iterator. */
|
|
it = PyObject_GetIter(seq);
|
|
if (it == NULL)
|
|
return NULL;
|
|
|
|
/* create filterfalseobject structure */
|
|
lz = (filterfalseobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(it);
|
|
return NULL;
|
|
}
|
|
lz->func = Py_NewRef(func);
|
|
lz->it = it;
|
|
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
filterfalse_dealloc(filterfalseobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->func);
|
|
Py_XDECREF(lz->it);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->it);
|
|
Py_VISIT(lz->func);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
filterfalse_next(filterfalseobject *lz)
|
|
{
|
|
PyObject *item;
|
|
PyObject *it = lz->it;
|
|
long ok;
|
|
PyObject *(*iternext)(PyObject *);
|
|
|
|
iternext = *Py_TYPE(it)->tp_iternext;
|
|
for (;;) {
|
|
item = iternext(it);
|
|
if (item == NULL)
|
|
return NULL;
|
|
|
|
if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
|
|
ok = PyObject_IsTrue(item);
|
|
} else {
|
|
PyObject *good;
|
|
good = PyObject_CallOneArg(lz->func, item);
|
|
if (good == NULL) {
|
|
Py_DECREF(item);
|
|
return NULL;
|
|
}
|
|
ok = PyObject_IsTrue(good);
|
|
Py_DECREF(good);
|
|
}
|
|
if (ok == 0)
|
|
return item;
|
|
Py_DECREF(item);
|
|
if (ok < 0)
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
static PyObject *
|
|
filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
|
|
}
|
|
|
|
static PyMethodDef filterfalse_methods[] = {
|
|
{"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot filterfalse_slots[] = {
|
|
{Py_tp_dealloc, filterfalse_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_filterfalse__doc__},
|
|
{Py_tp_traverse, filterfalse_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, filterfalse_next},
|
|
{Py_tp_methods, filterfalse_methods},
|
|
{Py_tp_new, itertools_filterfalse},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec filterfalse_spec = {
|
|
.name = "itertools.filterfalse",
|
|
.basicsize = sizeof(filterfalseobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = filterfalse_slots,
|
|
};
|
|
|
|
|
|
/* count object ************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
Py_ssize_t cnt;
|
|
PyObject *long_cnt;
|
|
PyObject *long_step;
|
|
} countobject;
|
|
|
|
/* Counting logic and invariants:
|
|
|
|
fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
|
|
|
|
assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
|
|
Advances with: cnt += 1
|
|
When count hits Y_SSIZE_T_MAX, switch to slow_mode.
|
|
|
|
slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
|
|
|
|
assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
|
|
All counting is done with python objects (no overflows or underflows).
|
|
Advances with: long_cnt += long_step
|
|
Step may be zero -- effectively a slow version of repeat(cnt).
|
|
Either long_cnt or long_step may be a float, Fraction, or Decimal.
|
|
*/
|
|
|
|
/*[clinic input]
|
|
@classmethod
|
|
itertools.count.__new__
|
|
start as long_cnt: object(c_default="NULL") = 0
|
|
step as long_step: object(c_default="NULL") = 1
|
|
Return a count object whose .__next__() method returns consecutive values.
|
|
|
|
Equivalent to:
|
|
def count(firstval=0, step=1):
|
|
x = firstval
|
|
while 1:
|
|
yield x
|
|
x += step
|
|
[clinic start generated code]*/
|
|
|
|
static PyObject *
|
|
itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
|
|
PyObject *long_step)
|
|
/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
|
|
{
|
|
countobject *lz;
|
|
int fast_mode;
|
|
Py_ssize_t cnt = 0;
|
|
long step;
|
|
|
|
if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
|
|
(long_step != NULL && !PyNumber_Check(long_step))) {
|
|
PyErr_SetString(PyExc_TypeError, "a number is required");
|
|
return NULL;
|
|
}
|
|
|
|
fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
|
|
(long_step == NULL || PyLong_Check(long_step));
|
|
|
|
/* If not specified, start defaults to 0 */
|
|
if (long_cnt != NULL) {
|
|
if (fast_mode) {
|
|
assert(PyLong_Check(long_cnt));
|
|
cnt = PyLong_AsSsize_t(long_cnt);
|
|
if (cnt == -1 && PyErr_Occurred()) {
|
|
PyErr_Clear();
|
|
fast_mode = 0;
|
|
}
|
|
}
|
|
} else {
|
|
cnt = 0;
|
|
long_cnt = _PyLong_GetZero();
|
|
}
|
|
Py_INCREF(long_cnt);
|
|
|
|
/* If not specified, step defaults to 1 */
|
|
if (long_step == NULL) {
|
|
long_step = _PyLong_GetOne();
|
|
}
|
|
Py_INCREF(long_step);
|
|
|
|
assert(long_cnt != NULL && long_step != NULL);
|
|
|
|
/* Fast mode only works when the step is 1 */
|
|
if (fast_mode) {
|
|
assert(PyLong_Check(long_step));
|
|
step = PyLong_AsLong(long_step);
|
|
if (step != 1) {
|
|
fast_mode = 0;
|
|
if (step == -1 && PyErr_Occurred())
|
|
PyErr_Clear();
|
|
}
|
|
}
|
|
|
|
if (fast_mode)
|
|
Py_CLEAR(long_cnt);
|
|
else
|
|
cnt = PY_SSIZE_T_MAX;
|
|
|
|
assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
|
|
(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
|
|
assert(!fast_mode ||
|
|
(PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
|
|
|
|
/* create countobject structure */
|
|
lz = (countobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_XDECREF(long_cnt);
|
|
Py_DECREF(long_step);
|
|
return NULL;
|
|
}
|
|
lz->cnt = cnt;
|
|
lz->long_cnt = long_cnt;
|
|
lz->long_step = long_step;
|
|
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
count_dealloc(countobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->long_cnt);
|
|
Py_XDECREF(lz->long_step);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
count_traverse(countobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->long_cnt);
|
|
Py_VISIT(lz->long_step);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
count_nextlong(countobject *lz)
|
|
{
|
|
PyObject *long_cnt;
|
|
PyObject *stepped_up;
|
|
|
|
long_cnt = lz->long_cnt;
|
|
if (long_cnt == NULL) {
|
|
/* Switch to slow_mode */
|
|
long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
|
|
if (long_cnt == NULL)
|
|
return NULL;
|
|
}
|
|
assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
|
|
|
|
stepped_up = PyNumber_Add(long_cnt, lz->long_step);
|
|
if (stepped_up == NULL)
|
|
return NULL;
|
|
lz->long_cnt = stepped_up;
|
|
return long_cnt;
|
|
}
|
|
|
|
static PyObject *
|
|
count_next(countobject *lz)
|
|
{
|
|
if (lz->cnt == PY_SSIZE_T_MAX)
|
|
return count_nextlong(lz);
|
|
return PyLong_FromSsize_t(lz->cnt++);
|
|
}
|
|
|
|
static PyObject *
|
|
count_repr(countobject *lz)
|
|
{
|
|
if (lz->cnt != PY_SSIZE_T_MAX)
|
|
return PyUnicode_FromFormat("%s(%zd)",
|
|
_PyType_Name(Py_TYPE(lz)), lz->cnt);
|
|
|
|
if (PyLong_Check(lz->long_step)) {
|
|
long step = PyLong_AsLong(lz->long_step);
|
|
if (step == -1 && PyErr_Occurred()) {
|
|
PyErr_Clear();
|
|
}
|
|
if (step == 1) {
|
|
/* Don't display step when it is an integer equal to 1 */
|
|
return PyUnicode_FromFormat("%s(%R)",
|
|
_PyType_Name(Py_TYPE(lz)),
|
|
lz->long_cnt);
|
|
}
|
|
}
|
|
return PyUnicode_FromFormat("%s(%R, %R)",
|
|
_PyType_Name(Py_TYPE(lz)),
|
|
lz->long_cnt, lz->long_step);
|
|
}
|
|
|
|
static PyObject *
|
|
count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
if (lz->cnt == PY_SSIZE_T_MAX)
|
|
return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
|
|
return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
|
|
}
|
|
|
|
static PyMethodDef count_methods[] = {
|
|
{"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
static PyType_Slot count_slots[] = {
|
|
{Py_tp_dealloc, count_dealloc},
|
|
{Py_tp_repr, count_repr},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)itertools_count__doc__},
|
|
{Py_tp_traverse, count_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, count_next},
|
|
{Py_tp_methods, count_methods},
|
|
{Py_tp_new, itertools_count},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec count_spec = {
|
|
.name = "itertools.count",
|
|
.basicsize = sizeof(countobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = count_slots,
|
|
};
|
|
|
|
|
|
/* repeat object ************************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
PyObject *element;
|
|
Py_ssize_t cnt;
|
|
} repeatobject;
|
|
|
|
static PyObject *
|
|
repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
|
{
|
|
repeatobject *ro;
|
|
PyObject *element;
|
|
Py_ssize_t cnt = -1, n_args;
|
|
static char *kwargs[] = {"object", "times", NULL};
|
|
|
|
n_args = PyTuple_GET_SIZE(args);
|
|
if (kwds != NULL)
|
|
n_args += PyDict_GET_SIZE(kwds);
|
|
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
|
|
&element, &cnt))
|
|
return NULL;
|
|
/* Does user supply times argument? */
|
|
if (n_args == 2 && cnt < 0)
|
|
cnt = 0;
|
|
|
|
ro = (repeatobject *)type->tp_alloc(type, 0);
|
|
if (ro == NULL)
|
|
return NULL;
|
|
ro->element = Py_NewRef(element);
|
|
ro->cnt = cnt;
|
|
return (PyObject *)ro;
|
|
}
|
|
|
|
static void
|
|
repeat_dealloc(repeatobject *ro)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(ro);
|
|
PyObject_GC_UnTrack(ro);
|
|
Py_XDECREF(ro->element);
|
|
tp->tp_free(ro);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(ro));
|
|
Py_VISIT(ro->element);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
repeat_next(repeatobject *ro)
|
|
{
|
|
if (ro->cnt == 0)
|
|
return NULL;
|
|
if (ro->cnt > 0)
|
|
ro->cnt--;
|
|
return Py_NewRef(ro->element);
|
|
}
|
|
|
|
static PyObject *
|
|
repeat_repr(repeatobject *ro)
|
|
{
|
|
if (ro->cnt == -1)
|
|
return PyUnicode_FromFormat("%s(%R)",
|
|
_PyType_Name(Py_TYPE(ro)), ro->element);
|
|
else
|
|
return PyUnicode_FromFormat("%s(%R, %zd)",
|
|
_PyType_Name(Py_TYPE(ro)), ro->element,
|
|
ro->cnt);
|
|
}
|
|
|
|
static PyObject *
|
|
repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
if (ro->cnt == -1) {
|
|
PyErr_SetString(PyExc_TypeError, "len() of unsized object");
|
|
return NULL;
|
|
}
|
|
return PyLong_FromSize_t(ro->cnt);
|
|
}
|
|
|
|
PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
|
|
|
|
static PyObject *
|
|
repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
/* unpickle this so that a new repeat iterator is constructed with an
|
|
* object, then call __setstate__ on it to set cnt
|
|
*/
|
|
if (ro->cnt >= 0)
|
|
return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
|
|
else
|
|
return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
|
|
}
|
|
|
|
static PyMethodDef repeat_methods[] = {
|
|
{"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
|
|
{"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
PyDoc_STRVAR(repeat_doc,
|
|
"repeat(object [,times]) -> create an iterator which returns the object\n\
|
|
for the specified number of times. If not specified, returns the object\n\
|
|
endlessly.");
|
|
|
|
static PyType_Slot repeat_slots[] = {
|
|
{Py_tp_dealloc, repeat_dealloc},
|
|
{Py_tp_repr, repeat_repr},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)repeat_doc},
|
|
{Py_tp_traverse, repeat_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, repeat_next},
|
|
{Py_tp_methods, repeat_methods},
|
|
{Py_tp_new, repeat_new},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec repeat_spec = {
|
|
.name = "itertools.repeat",
|
|
.basicsize = sizeof(repeatobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = repeat_slots,
|
|
};
|
|
|
|
|
|
/* ziplongest object *********************************************************/
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
Py_ssize_t tuplesize;
|
|
Py_ssize_t numactive;
|
|
PyObject *ittuple; /* tuple of iterators */
|
|
PyObject *result;
|
|
PyObject *fillvalue;
|
|
} ziplongestobject;
|
|
|
|
static PyObject *
|
|
zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
|
{
|
|
ziplongestobject *lz;
|
|
Py_ssize_t i;
|
|
PyObject *ittuple; /* tuple of iterators */
|
|
PyObject *result;
|
|
PyObject *fillvalue = Py_None;
|
|
Py_ssize_t tuplesize;
|
|
|
|
if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
|
|
fillvalue = NULL;
|
|
if (PyDict_GET_SIZE(kwds) == 1) {
|
|
fillvalue = PyDict_GetItemWithError(kwds, &_Py_ID(fillvalue));
|
|
}
|
|
if (fillvalue == NULL) {
|
|
if (!PyErr_Occurred()) {
|
|
PyErr_SetString(PyExc_TypeError,
|
|
"zip_longest() got an unexpected keyword argument");
|
|
}
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
/* args must be a tuple */
|
|
assert(PyTuple_Check(args));
|
|
tuplesize = PyTuple_GET_SIZE(args);
|
|
|
|
/* obtain iterators */
|
|
ittuple = PyTuple_New(tuplesize);
|
|
if (ittuple == NULL)
|
|
return NULL;
|
|
for (i=0; i < tuplesize; i++) {
|
|
PyObject *item = PyTuple_GET_ITEM(args, i);
|
|
PyObject *it = PyObject_GetIter(item);
|
|
if (it == NULL) {
|
|
Py_DECREF(ittuple);
|
|
return NULL;
|
|
}
|
|
PyTuple_SET_ITEM(ittuple, i, it);
|
|
}
|
|
|
|
/* create a result holder */
|
|
result = PyTuple_New(tuplesize);
|
|
if (result == NULL) {
|
|
Py_DECREF(ittuple);
|
|
return NULL;
|
|
}
|
|
for (i=0 ; i < tuplesize ; i++) {
|
|
Py_INCREF(Py_None);
|
|
PyTuple_SET_ITEM(result, i, Py_None);
|
|
}
|
|
|
|
/* create ziplongestobject structure */
|
|
lz = (ziplongestobject *)type->tp_alloc(type, 0);
|
|
if (lz == NULL) {
|
|
Py_DECREF(ittuple);
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
}
|
|
lz->ittuple = ittuple;
|
|
lz->tuplesize = tuplesize;
|
|
lz->numactive = tuplesize;
|
|
lz->result = result;
|
|
lz->fillvalue = Py_NewRef(fillvalue);
|
|
return (PyObject *)lz;
|
|
}
|
|
|
|
static void
|
|
zip_longest_dealloc(ziplongestobject *lz)
|
|
{
|
|
PyTypeObject *tp = Py_TYPE(lz);
|
|
PyObject_GC_UnTrack(lz);
|
|
Py_XDECREF(lz->ittuple);
|
|
Py_XDECREF(lz->result);
|
|
Py_XDECREF(lz->fillvalue);
|
|
tp->tp_free(lz);
|
|
Py_DECREF(tp);
|
|
}
|
|
|
|
static int
|
|
zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
|
|
{
|
|
Py_VISIT(Py_TYPE(lz));
|
|
Py_VISIT(lz->ittuple);
|
|
Py_VISIT(lz->result);
|
|
Py_VISIT(lz->fillvalue);
|
|
return 0;
|
|
}
|
|
|
|
static PyObject *
|
|
zip_longest_next(ziplongestobject *lz)
|
|
{
|
|
Py_ssize_t i;
|
|
Py_ssize_t tuplesize = lz->tuplesize;
|
|
PyObject *result = lz->result;
|
|
PyObject *it;
|
|
PyObject *item;
|
|
PyObject *olditem;
|
|
|
|
if (tuplesize == 0)
|
|
return NULL;
|
|
if (lz->numactive == 0)
|
|
return NULL;
|
|
if (Py_REFCNT(result) == 1) {
|
|
Py_INCREF(result);
|
|
for (i=0 ; i < tuplesize ; i++) {
|
|
it = PyTuple_GET_ITEM(lz->ittuple, i);
|
|
if (it == NULL) {
|
|
item = Py_NewRef(lz->fillvalue);
|
|
} else {
|
|
item = PyIter_Next(it);
|
|
if (item == NULL) {
|
|
lz->numactive -= 1;
|
|
if (lz->numactive == 0 || PyErr_Occurred()) {
|
|
lz->numactive = 0;
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
} else {
|
|
item = Py_NewRef(lz->fillvalue);
|
|
PyTuple_SET_ITEM(lz->ittuple, i, NULL);
|
|
Py_DECREF(it);
|
|
}
|
|
}
|
|
}
|
|
olditem = PyTuple_GET_ITEM(result, i);
|
|
PyTuple_SET_ITEM(result, i, item);
|
|
Py_DECREF(olditem);
|
|
}
|
|
// bpo-42536: The GC may have untracked this result tuple. Since we're
|
|
// recycling it, make sure it's tracked again:
|
|
if (!_PyObject_GC_IS_TRACKED(result)) {
|
|
_PyObject_GC_TRACK(result);
|
|
}
|
|
} else {
|
|
result = PyTuple_New(tuplesize);
|
|
if (result == NULL)
|
|
return NULL;
|
|
for (i=0 ; i < tuplesize ; i++) {
|
|
it = PyTuple_GET_ITEM(lz->ittuple, i);
|
|
if (it == NULL) {
|
|
item = Py_NewRef(lz->fillvalue);
|
|
} else {
|
|
item = PyIter_Next(it);
|
|
if (item == NULL) {
|
|
lz->numactive -= 1;
|
|
if (lz->numactive == 0 || PyErr_Occurred()) {
|
|
lz->numactive = 0;
|
|
Py_DECREF(result);
|
|
return NULL;
|
|
} else {
|
|
item = Py_NewRef(lz->fillvalue);
|
|
PyTuple_SET_ITEM(lz->ittuple, i, NULL);
|
|
Py_DECREF(it);
|
|
}
|
|
}
|
|
}
|
|
PyTuple_SET_ITEM(result, i, item);
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
static PyObject *
|
|
zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
|
|
{
|
|
|
|
/* Create a new tuple with empty sequences where appropriate to pickle.
|
|
* Then use setstate to set the fillvalue
|
|
*/
|
|
int i;
|
|
PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
|
|
|
|
if (args == NULL)
|
|
return NULL;
|
|
for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
|
|
PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
|
|
if (elem == NULL) {
|
|
elem = PyTuple_New(0);
|
|
if (elem == NULL) {
|
|
Py_DECREF(args);
|
|
return NULL;
|
|
}
|
|
} else
|
|
Py_INCREF(elem);
|
|
PyTuple_SET_ITEM(args, i, elem);
|
|
}
|
|
return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
|
|
}
|
|
|
|
static PyObject *
|
|
zip_longest_setstate(ziplongestobject *lz, PyObject *state)
|
|
{
|
|
Py_INCREF(state);
|
|
Py_XSETREF(lz->fillvalue, state);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyMethodDef zip_longest_methods[] = {
|
|
{"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
|
|
reduce_doc},
|
|
{"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
|
|
setstate_doc},
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
PyDoc_STRVAR(zip_longest_doc,
|
|
"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
|
|
\n\
|
|
Return a zip_longest object whose .__next__() method returns a tuple where\n\
|
|
the i-th element comes from the i-th iterable argument. The .__next__()\n\
|
|
method continues until the longest iterable in the argument sequence\n\
|
|
is exhausted and then it raises StopIteration. When the shorter iterables\n\
|
|
are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
|
|
defaults to None or can be specified by a keyword argument.\n\
|
|
");
|
|
|
|
static PyType_Slot ziplongest_slots[] = {
|
|
{Py_tp_dealloc, zip_longest_dealloc},
|
|
{Py_tp_getattro, PyObject_GenericGetAttr},
|
|
{Py_tp_doc, (void *)zip_longest_doc},
|
|
{Py_tp_traverse, zip_longest_traverse},
|
|
{Py_tp_iter, PyObject_SelfIter},
|
|
{Py_tp_iternext, zip_longest_next},
|
|
{Py_tp_methods, zip_longest_methods},
|
|
{Py_tp_new, zip_longest_new},
|
|
{Py_tp_free, PyObject_GC_Del},
|
|
{0, NULL},
|
|
};
|
|
|
|
static PyType_Spec ziplongest_spec = {
|
|
.name = "itertools.zip_longest",
|
|
.basicsize = sizeof(ziplongestobject),
|
|
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
|
|
Py_TPFLAGS_IMMUTABLETYPE),
|
|
.slots = ziplongest_slots,
|
|
};
|
|
|
|
|
|
/* module level code ********************************************************/
|
|
|
|
PyDoc_STRVAR(module_doc,
|
|
"Functional tools for creating and using iterators.\n\
|
|
\n\
|
|
Infinite iterators:\n\
|
|
count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
|
|
cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
|
|
repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
|
|
\n\
|
|
Iterators terminating on the shortest input sequence:\n\
|
|
accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
|
|
batched(p, n) --> [p0, p1, ..., p_n-1], [p_n, p_n+1, ..., p_2n-1], ...\n\
|
|
chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
|
|
chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
|
|
compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
|
|
dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
|
|
groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
|
|
filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
|
|
islice(seq, [start,] stop [, step]) --> elements from\n\
|
|
seq[start:stop:step]\n\
|
|
pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
|
|
starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
|
|
tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
|
|
takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
|
|
zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
|
|
\n\
|
|
Combinatoric generators:\n\
|
|
product(p, q, ... [repeat=1]) --> cartesian product\n\
|
|
permutations(p[, r])\n\
|
|
combinations(p, r)\n\
|
|
combinations_with_replacement(p, r)\n\
|
|
");
|
|
|
|
static int
|
|
itertoolsmodule_traverse(PyObject *mod, visitproc visit, void *arg)
|
|
{
|
|
itertools_state *state = get_module_state(mod);
|
|
Py_VISIT(state->accumulate_type);
|
|
Py_VISIT(state->batched_type);
|
|
Py_VISIT(state->chain_type);
|
|
Py_VISIT(state->combinations_type);
|
|
Py_VISIT(state->compress_type);
|
|
Py_VISIT(state->count_type);
|
|
Py_VISIT(state->cwr_type);
|
|
Py_VISIT(state->cycle_type);
|
|
Py_VISIT(state->dropwhile_type);
|
|
Py_VISIT(state->filterfalse_type);
|
|
Py_VISIT(state->groupby_type);
|
|
Py_VISIT(state->_grouper_type);
|
|
Py_VISIT(state->islice_type);
|
|
Py_VISIT(state->pairwise_type);
|
|
Py_VISIT(state->permutations_type);
|
|
Py_VISIT(state->product_type);
|
|
Py_VISIT(state->repeat_type);
|
|
Py_VISIT(state->starmap_type);
|
|
Py_VISIT(state->takewhile_type);
|
|
Py_VISIT(state->tee_type);
|
|
Py_VISIT(state->teedataobject_type);
|
|
Py_VISIT(state->ziplongest_type);
|
|
return 0;
|
|
}
|
|
|
|
static int
|
|
itertoolsmodule_clear(PyObject *mod)
|
|
{
|
|
itertools_state *state = get_module_state(mod);
|
|
Py_CLEAR(state->accumulate_type);
|
|
Py_CLEAR(state->batched_type);
|
|
Py_CLEAR(state->chain_type);
|
|
Py_CLEAR(state->combinations_type);
|
|
Py_CLEAR(state->compress_type);
|
|
Py_CLEAR(state->count_type);
|
|
Py_CLEAR(state->cwr_type);
|
|
Py_CLEAR(state->cycle_type);
|
|
Py_CLEAR(state->dropwhile_type);
|
|
Py_CLEAR(state->filterfalse_type);
|
|
Py_CLEAR(state->groupby_type);
|
|
Py_CLEAR(state->_grouper_type);
|
|
Py_CLEAR(state->islice_type);
|
|
Py_CLEAR(state->pairwise_type);
|
|
Py_CLEAR(state->permutations_type);
|
|
Py_CLEAR(state->product_type);
|
|
Py_CLEAR(state->repeat_type);
|
|
Py_CLEAR(state->starmap_type);
|
|
Py_CLEAR(state->takewhile_type);
|
|
Py_CLEAR(state->tee_type);
|
|
Py_CLEAR(state->teedataobject_type);
|
|
Py_CLEAR(state->ziplongest_type);
|
|
return 0;
|
|
}
|
|
|
|
static void
|
|
itertoolsmodule_free(void *mod)
|
|
{
|
|
(void)itertoolsmodule_clear((PyObject *)mod);
|
|
}
|
|
|
|
#define ADD_TYPE(module, type, spec) \
|
|
do { \
|
|
type = (PyTypeObject *)PyType_FromModuleAndSpec(module, spec, NULL); \
|
|
if (type == NULL) { \
|
|
return -1; \
|
|
} \
|
|
if (PyModule_AddType(module, type) < 0) { \
|
|
return -1; \
|
|
} \
|
|
} while (0)
|
|
|
|
static int
|
|
itertoolsmodule_exec(PyObject *mod)
|
|
{
|
|
itertools_state *state = get_module_state(mod);
|
|
ADD_TYPE(mod, state->accumulate_type, &accumulate_spec);
|
|
ADD_TYPE(mod, state->batched_type, &batched_spec);
|
|
ADD_TYPE(mod, state->chain_type, &chain_spec);
|
|
ADD_TYPE(mod, state->combinations_type, &combinations_spec);
|
|
ADD_TYPE(mod, state->compress_type, &compress_spec);
|
|
ADD_TYPE(mod, state->count_type, &count_spec);
|
|
ADD_TYPE(mod, state->cwr_type, &cwr_spec);
|
|
ADD_TYPE(mod, state->cycle_type, &cycle_spec);
|
|
ADD_TYPE(mod, state->dropwhile_type, &dropwhile_spec);
|
|
ADD_TYPE(mod, state->filterfalse_type, &filterfalse_spec);
|
|
ADD_TYPE(mod, state->groupby_type, &groupby_spec);
|
|
ADD_TYPE(mod, state->_grouper_type, &_grouper_spec);
|
|
ADD_TYPE(mod, state->islice_type, &islice_spec);
|
|
ADD_TYPE(mod, state->pairwise_type, &pairwise_spec);
|
|
ADD_TYPE(mod, state->permutations_type, &permutations_spec);
|
|
ADD_TYPE(mod, state->product_type, &product_spec);
|
|
ADD_TYPE(mod, state->repeat_type, &repeat_spec);
|
|
ADD_TYPE(mod, state->starmap_type, &starmap_spec);
|
|
ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
|
|
ADD_TYPE(mod, state->tee_type, &tee_spec);
|
|
ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
|
|
ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
|
|
|
|
Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
|
|
return 0;
|
|
}
|
|
|
|
static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
|
|
{Py_mod_exec, itertoolsmodule_exec},
|
|
{0, NULL}
|
|
};
|
|
|
|
static PyMethodDef module_methods[] = {
|
|
ITERTOOLS_TEE_METHODDEF
|
|
{NULL, NULL} /* sentinel */
|
|
};
|
|
|
|
|
|
static struct PyModuleDef itertoolsmodule = {
|
|
.m_base = PyModuleDef_HEAD_INIT,
|
|
.m_name = "itertools",
|
|
.m_doc = module_doc,
|
|
.m_size = sizeof(itertools_state),
|
|
.m_methods = module_methods,
|
|
.m_slots = itertoolsmodule_slots,
|
|
.m_traverse = itertoolsmodule_traverse,
|
|
.m_clear = itertoolsmodule_clear,
|
|
.m_free = itertoolsmodule_free,
|
|
};
|
|
|
|
PyMODINIT_FUNC
|
|
PyInit_itertools(void)
|
|
{
|
|
return PyModuleDef_Init(&itertoolsmodule);
|
|
}
|