ยปCore Development>Code coverage>Modules/itertoolsmodule.c

Python code coverage for Modules/itertoolsmodule.c

#countcontent
1n/a
2n/a#define PY_SSIZE_T_CLEAN
3n/a#include "Python.h"
4n/a#include "structmember.h"
5n/a
6n/a/* Itertools module written and maintained
7n/a by Raymond D. Hettinger <python@rcn.com>
8n/a*/
9n/a
10n/a
11n/a/* groupby object ************************************************************/
12n/a
13n/atypedef struct {
14n/a PyObject_HEAD
15n/a PyObject *it;
16n/a PyObject *keyfunc;
17n/a PyObject *tgtkey;
18n/a PyObject *currkey;
19n/a PyObject *currvalue;
20n/a} groupbyobject;
21n/a
22n/astatic PyTypeObject groupby_type;
23n/astatic PyObject *_grouper_create(groupbyobject *, PyObject *);
24n/a
25n/astatic PyObject *
26n/agroupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
27n/a{
28n/a static char *kwargs[] = {"iterable", "key", NULL};
29n/a groupbyobject *gbo;
30n/a PyObject *it, *keyfunc = Py_None;
31n/a
32n/a if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
33n/a &it, &keyfunc))
34n/a return NULL;
35n/a
36n/a gbo = (groupbyobject *)type->tp_alloc(type, 0);
37n/a if (gbo == NULL)
38n/a return NULL;
39n/a gbo->tgtkey = NULL;
40n/a gbo->currkey = NULL;
41n/a gbo->currvalue = NULL;
42n/a gbo->keyfunc = keyfunc;
43n/a Py_INCREF(keyfunc);
44n/a gbo->it = PyObject_GetIter(it);
45n/a if (gbo->it == NULL) {
46n/a Py_DECREF(gbo);
47n/a return NULL;
48n/a }
49n/a return (PyObject *)gbo;
50n/a}
51n/a
52n/astatic void
53n/agroupby_dealloc(groupbyobject *gbo)
54n/a{
55n/a PyObject_GC_UnTrack(gbo);
56n/a Py_XDECREF(gbo->it);
57n/a Py_XDECREF(gbo->keyfunc);
58n/a Py_XDECREF(gbo->tgtkey);
59n/a Py_XDECREF(gbo->currkey);
60n/a Py_XDECREF(gbo->currvalue);
61n/a Py_TYPE(gbo)->tp_free(gbo);
62n/a}
63n/a
64n/astatic int
65n/agroupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
66n/a{
67n/a Py_VISIT(gbo->it);
68n/a Py_VISIT(gbo->keyfunc);
69n/a Py_VISIT(gbo->tgtkey);
70n/a Py_VISIT(gbo->currkey);
71n/a Py_VISIT(gbo->currvalue);
72n/a return 0;
73n/a}
74n/a
75n/astatic PyObject *
76n/agroupby_next(groupbyobject *gbo)
77n/a{
78n/a PyObject *newvalue, *newkey, *r, *grouper;
79n/a
80n/a /* skip to next iteration group */
81n/a for (;;) {
82n/a if (gbo->currkey == NULL)
83n/a /* pass */;
84n/a else if (gbo->tgtkey == NULL)
85n/a break;
86n/a else {
87n/a int rcmp;
88n/a
89n/a rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
90n/a if (rcmp == -1)
91n/a return NULL;
92n/a else if (rcmp == 0)
93n/a break;
94n/a }
95n/a
96n/a newvalue = PyIter_Next(gbo->it);
97n/a if (newvalue == NULL)
98n/a return NULL;
99n/a
100n/a if (gbo->keyfunc == Py_None) {
101n/a newkey = newvalue;
102n/a Py_INCREF(newvalue);
103n/a } else {
104n/a newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
105n/a if (newkey == NULL) {
106n/a Py_DECREF(newvalue);
107n/a return NULL;
108n/a }
109n/a }
110n/a
111n/a Py_XSETREF(gbo->currkey, newkey);
112n/a Py_XSETREF(gbo->currvalue, newvalue);
113n/a }
114n/a
115n/a Py_INCREF(gbo->currkey);
116n/a Py_XSETREF(gbo->tgtkey, gbo->currkey);
117n/a
118n/a grouper = _grouper_create(gbo, gbo->tgtkey);
119n/a if (grouper == NULL)
120n/a return NULL;
121n/a
122n/a r = PyTuple_Pack(2, gbo->currkey, grouper);
123n/a Py_DECREF(grouper);
124n/a return r;
125n/a}
126n/a
127n/astatic PyObject *
128n/agroupby_reduce(groupbyobject *lz)
129n/a{
130n/a /* reduce as a 'new' call with an optional 'setstate' if groupby
131n/a * has started
132n/a */
133n/a PyObject *value;
134n/a if (lz->tgtkey && lz->currkey && lz->currvalue)
135n/a value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
136n/a lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
137n/a else
138n/a value = Py_BuildValue("O(OO)", Py_TYPE(lz),
139n/a lz->it, lz->keyfunc);
140n/a
141n/a return value;
142n/a}
143n/a
144n/aPyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
145n/a
146n/astatic PyObject *
147n/agroupby_setstate(groupbyobject *lz, PyObject *state)
148n/a{
149n/a PyObject *currkey, *currvalue, *tgtkey;
150n/a if (!PyTuple_Check(state)) {
151n/a PyErr_SetString(PyExc_TypeError, "state is not a tuple");
152n/a return NULL;
153n/a }
154n/a if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
155n/a return NULL;
156n/a }
157n/a Py_INCREF(currkey);
158n/a Py_XSETREF(lz->currkey, currkey);
159n/a Py_INCREF(currvalue);
160n/a Py_XSETREF(lz->currvalue, currvalue);
161n/a Py_INCREF(tgtkey);
162n/a Py_XSETREF(lz->tgtkey, tgtkey);
163n/a Py_RETURN_NONE;
164n/a}
165n/a
166n/aPyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
167n/a
168n/astatic PyMethodDef groupby_methods[] = {
169n/a {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
170n/a reduce_doc},
171n/a {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
172n/a setstate_doc},
173n/a {NULL, NULL} /* sentinel */
174n/a};
175n/a
176n/aPyDoc_STRVAR(groupby_doc,
177n/a"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
178n/a(key, sub-iterator) grouped by each value of key(value).\n");
179n/a
180n/astatic PyTypeObject groupby_type = {
181n/a PyVarObject_HEAD_INIT(NULL, 0)
182n/a "itertools.groupby", /* tp_name */
183n/a sizeof(groupbyobject), /* tp_basicsize */
184n/a 0, /* tp_itemsize */
185n/a /* methods */
186n/a (destructor)groupby_dealloc, /* tp_dealloc */
187n/a 0, /* tp_print */
188n/a 0, /* tp_getattr */
189n/a 0, /* tp_setattr */
190n/a 0, /* tp_reserved */
191n/a 0, /* tp_repr */
192n/a 0, /* tp_as_number */
193n/a 0, /* tp_as_sequence */
194n/a 0, /* tp_as_mapping */
195n/a 0, /* tp_hash */
196n/a 0, /* tp_call */
197n/a 0, /* tp_str */
198n/a PyObject_GenericGetAttr, /* tp_getattro */
199n/a 0, /* tp_setattro */
200n/a 0, /* tp_as_buffer */
201n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
202n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
203n/a groupby_doc, /* tp_doc */
204n/a (traverseproc)groupby_traverse, /* tp_traverse */
205n/a 0, /* tp_clear */
206n/a 0, /* tp_richcompare */
207n/a 0, /* tp_weaklistoffset */
208n/a PyObject_SelfIter, /* tp_iter */
209n/a (iternextfunc)groupby_next, /* tp_iternext */
210n/a groupby_methods, /* tp_methods */
211n/a 0, /* tp_members */
212n/a 0, /* tp_getset */
213n/a 0, /* tp_base */
214n/a 0, /* tp_dict */
215n/a 0, /* tp_descr_get */
216n/a 0, /* tp_descr_set */
217n/a 0, /* tp_dictoffset */
218n/a 0, /* tp_init */
219n/a 0, /* tp_alloc */
220n/a groupby_new, /* tp_new */
221n/a PyObject_GC_Del, /* tp_free */
222n/a};
223n/a
224n/a
225n/a/* _grouper object (internal) ************************************************/
226n/a
227n/atypedef struct {
228n/a PyObject_HEAD
229n/a PyObject *parent;
230n/a PyObject *tgtkey;
231n/a} _grouperobject;
232n/a
233n/astatic PyTypeObject _grouper_type;
234n/a
235n/astatic PyObject *
236n/a_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
237n/a{
238n/a PyObject *parent, *tgtkey;
239n/a
240n/a if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
241n/a return NULL;
242n/a
243n/a return _grouper_create((groupbyobject*) parent, tgtkey);
244n/a}
245n/a
246n/astatic PyObject *
247n/a_grouper_create(groupbyobject *parent, PyObject *tgtkey)
248n/a{
249n/a _grouperobject *igo;
250n/a
251n/a igo = PyObject_GC_New(_grouperobject, &_grouper_type);
252n/a if (igo == NULL)
253n/a return NULL;
254n/a igo->parent = (PyObject *)parent;
255n/a Py_INCREF(parent);
256n/a igo->tgtkey = tgtkey;
257n/a Py_INCREF(tgtkey);
258n/a
259n/a PyObject_GC_Track(igo);
260n/a return (PyObject *)igo;
261n/a}
262n/a
263n/astatic void
264n/a_grouper_dealloc(_grouperobject *igo)
265n/a{
266n/a PyObject_GC_UnTrack(igo);
267n/a Py_DECREF(igo->parent);
268n/a Py_DECREF(igo->tgtkey);
269n/a PyObject_GC_Del(igo);
270n/a}
271n/a
272n/astatic int
273n/a_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
274n/a{
275n/a Py_VISIT(igo->parent);
276n/a Py_VISIT(igo->tgtkey);
277n/a return 0;
278n/a}
279n/a
280n/astatic PyObject *
281n/a_grouper_next(_grouperobject *igo)
282n/a{
283n/a groupbyobject *gbo = (groupbyobject *)igo->parent;
284n/a PyObject *newvalue, *newkey, *r;
285n/a int rcmp;
286n/a
287n/a if (gbo->currvalue == NULL) {
288n/a newvalue = PyIter_Next(gbo->it);
289n/a if (newvalue == NULL)
290n/a return NULL;
291n/a
292n/a if (gbo->keyfunc == Py_None) {
293n/a newkey = newvalue;
294n/a Py_INCREF(newvalue);
295n/a } else {
296n/a newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
297n/a if (newkey == NULL) {
298n/a Py_DECREF(newvalue);
299n/a return NULL;
300n/a }
301n/a }
302n/a
303n/a assert(gbo->currkey == NULL);
304n/a gbo->currkey = newkey;
305n/a gbo->currvalue = newvalue;
306n/a }
307n/a
308n/a assert(gbo->currkey != NULL);
309n/a rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
310n/a if (rcmp <= 0)
311n/a /* got any error or current group is end */
312n/a return NULL;
313n/a
314n/a r = gbo->currvalue;
315n/a gbo->currvalue = NULL;
316n/a Py_CLEAR(gbo->currkey);
317n/a
318n/a return r;
319n/a}
320n/a
321n/astatic PyObject *
322n/a_grouper_reduce(_grouperobject *lz)
323n/a{
324n/a return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
325n/a}
326n/a
327n/astatic PyMethodDef _grouper_methods[] = {
328n/a {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
329n/a reduce_doc},
330n/a {NULL, NULL} /* sentinel */
331n/a};
332n/a
333n/a
334n/astatic PyTypeObject _grouper_type = {
335n/a PyVarObject_HEAD_INIT(NULL, 0)
336n/a "itertools._grouper", /* tp_name */
337n/a sizeof(_grouperobject), /* tp_basicsize */
338n/a 0, /* tp_itemsize */
339n/a /* methods */
340n/a (destructor)_grouper_dealloc, /* tp_dealloc */
341n/a 0, /* tp_print */
342n/a 0, /* tp_getattr */
343n/a 0, /* tp_setattr */
344n/a 0, /* tp_reserved */
345n/a 0, /* tp_repr */
346n/a 0, /* tp_as_number */
347n/a 0, /* tp_as_sequence */
348n/a 0, /* tp_as_mapping */
349n/a 0, /* tp_hash */
350n/a 0, /* tp_call */
351n/a 0, /* tp_str */
352n/a PyObject_GenericGetAttr, /* tp_getattro */
353n/a 0, /* tp_setattro */
354n/a 0, /* tp_as_buffer */
355n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
356n/a 0, /* tp_doc */
357n/a (traverseproc)_grouper_traverse, /* tp_traverse */
358n/a 0, /* tp_clear */
359n/a 0, /* tp_richcompare */
360n/a 0, /* tp_weaklistoffset */
361n/a PyObject_SelfIter, /* tp_iter */
362n/a (iternextfunc)_grouper_next, /* tp_iternext */
363n/a _grouper_methods, /* tp_methods */
364n/a 0, /* tp_members */
365n/a 0, /* tp_getset */
366n/a 0, /* tp_base */
367n/a 0, /* tp_dict */
368n/a 0, /* tp_descr_get */
369n/a 0, /* tp_descr_set */
370n/a 0, /* tp_dictoffset */
371n/a 0, /* tp_init */
372n/a 0, /* tp_alloc */
373n/a _grouper_new, /* tp_new */
374n/a PyObject_GC_Del, /* tp_free */
375n/a};
376n/a
377n/a
378n/a/* tee object and with supporting function and objects ***********************/
379n/a
380n/a/* The teedataobject pre-allocates space for LINKCELLS number of objects.
381n/a To help the object fit neatly inside cache lines (space for 16 to 32
382n/a pointers), the value should be a multiple of 16 minus space for
383n/a the other structure members including PyHEAD overhead. The larger the
384n/a value, the less memory overhead per object and the less time spent
385n/a allocating/deallocating new links. The smaller the number, the less
386n/a wasted space and the more rapid freeing of older data.
387n/a*/
388n/a#define LINKCELLS 57
389n/a
390n/atypedef struct {
391n/a PyObject_HEAD
392n/a PyObject *it;
393n/a int numread; /* 0 <= numread <= LINKCELLS */
394n/a PyObject *nextlink;
395n/a PyObject *(values[LINKCELLS]);
396n/a} teedataobject;
397n/a
398n/atypedef struct {
399n/a PyObject_HEAD
400n/a teedataobject *dataobj;
401n/a int index; /* 0 <= index <= LINKCELLS */
402n/a PyObject *weakreflist;
403n/a} teeobject;
404n/a
405n/astatic PyTypeObject teedataobject_type;
406n/a
407n/astatic PyObject *
408n/ateedataobject_newinternal(PyObject *it)
409n/a{
410n/a teedataobject *tdo;
411n/a
412n/a tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
413n/a if (tdo == NULL)
414n/a return NULL;
415n/a
416n/a tdo->numread = 0;
417n/a tdo->nextlink = NULL;
418n/a Py_INCREF(it);
419n/a tdo->it = it;
420n/a PyObject_GC_Track(tdo);
421n/a return (PyObject *)tdo;
422n/a}
423n/a
424n/astatic PyObject *
425n/ateedataobject_jumplink(teedataobject *tdo)
426n/a{
427n/a if (tdo->nextlink == NULL)
428n/a tdo->nextlink = teedataobject_newinternal(tdo->it);
429n/a Py_XINCREF(tdo->nextlink);
430n/a return tdo->nextlink;
431n/a}
432n/a
433n/astatic PyObject *
434n/ateedataobject_getitem(teedataobject *tdo, int i)
435n/a{
436n/a PyObject *value;
437n/a
438n/a assert(i < LINKCELLS);
439n/a if (i < tdo->numread)
440n/a value = tdo->values[i];
441n/a else {
442n/a /* this is the lead iterator, so fetch more data */
443n/a assert(i == tdo->numread);
444n/a value = PyIter_Next(tdo->it);
445n/a if (value == NULL)
446n/a return NULL;
447n/a tdo->numread++;
448n/a tdo->values[i] = value;
449n/a }
450n/a Py_INCREF(value);
451n/a return value;
452n/a}
453n/a
454n/astatic int
455n/ateedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
456n/a{
457n/a int i;
458n/a
459n/a Py_VISIT(tdo->it);
460n/a for (i = 0; i < tdo->numread; i++)
461n/a Py_VISIT(tdo->values[i]);
462n/a Py_VISIT(tdo->nextlink);
463n/a return 0;
464n/a}
465n/a
466n/astatic void
467n/ateedataobject_safe_decref(PyObject *obj)
468n/a{
469n/a while (obj && Py_TYPE(obj) == &teedataobject_type &&
470n/a Py_REFCNT(obj) == 1) {
471n/a PyObject *nextlink = ((teedataobject *)obj)->nextlink;
472n/a ((teedataobject *)obj)->nextlink = NULL;
473n/a Py_DECREF(obj);
474n/a obj = nextlink;
475n/a }
476n/a Py_XDECREF(obj);
477n/a}
478n/a
479n/astatic int
480n/ateedataobject_clear(teedataobject *tdo)
481n/a{
482n/a int i;
483n/a PyObject *tmp;
484n/a
485n/a Py_CLEAR(tdo->it);
486n/a for (i=0 ; i<tdo->numread ; i++)
487n/a Py_CLEAR(tdo->values[i]);
488n/a tmp = tdo->nextlink;
489n/a tdo->nextlink = NULL;
490n/a teedataobject_safe_decref(tmp);
491n/a return 0;
492n/a}
493n/a
494n/astatic void
495n/ateedataobject_dealloc(teedataobject *tdo)
496n/a{
497n/a PyObject_GC_UnTrack(tdo);
498n/a teedataobject_clear(tdo);
499n/a PyObject_GC_Del(tdo);
500n/a}
501n/a
502n/astatic PyObject *
503n/ateedataobject_reduce(teedataobject *tdo)
504n/a{
505n/a int i;
506n/a /* create a temporary list of already iterated values */
507n/a PyObject *values = PyList_New(tdo->numread);
508n/a
509n/a if (!values)
510n/a return NULL;
511n/a for (i=0 ; i<tdo->numread ; i++) {
512n/a Py_INCREF(tdo->values[i]);
513n/a PyList_SET_ITEM(values, i, tdo->values[i]);
514n/a }
515n/a return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
516n/a values,
517n/a tdo->nextlink ? tdo->nextlink : Py_None);
518n/a}
519n/a
520n/astatic PyTypeObject teedataobject_type;
521n/a
522n/astatic PyObject *
523n/ateedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
524n/a{
525n/a teedataobject *tdo;
526n/a PyObject *it, *values, *next;
527n/a Py_ssize_t i, len;
528n/a
529n/a assert(type == &teedataobject_type);
530n/a if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
531n/a return NULL;
532n/a
533n/a tdo = (teedataobject *)teedataobject_newinternal(it);
534n/a if (!tdo)
535n/a return NULL;
536n/a
537n/a len = PyList_GET_SIZE(values);
538n/a if (len > LINKCELLS)
539n/a goto err;
540n/a for (i=0; i<len; i++) {
541n/a tdo->values[i] = PyList_GET_ITEM(values, i);
542n/a Py_INCREF(tdo->values[i]);
543n/a }
544n/a /* len <= LINKCELLS < INT_MAX */
545n/a tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
546n/a
547n/a if (len == LINKCELLS) {
548n/a if (next != Py_None) {
549n/a if (Py_TYPE(next) != &teedataobject_type)
550n/a goto err;
551n/a assert(tdo->nextlink == NULL);
552n/a Py_INCREF(next);
553n/a tdo->nextlink = next;
554n/a }
555n/a } else {
556n/a if (next != Py_None)
557n/a goto err; /* shouldn't have a next if we are not full */
558n/a }
559n/a return (PyObject*)tdo;
560n/a
561n/aerr:
562n/a Py_XDECREF(tdo);
563n/a PyErr_SetString(PyExc_ValueError, "Invalid arguments");
564n/a return NULL;
565n/a}
566n/a
567n/astatic PyMethodDef teedataobject_methods[] = {
568n/a {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
569n/a reduce_doc},
570n/a {NULL, NULL} /* sentinel */
571n/a};
572n/a
573n/aPyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
574n/a
575n/astatic PyTypeObject teedataobject_type = {
576n/a PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
577n/a "itertools._tee_dataobject", /* tp_name */
578n/a sizeof(teedataobject), /* tp_basicsize */
579n/a 0, /* tp_itemsize */
580n/a /* methods */
581n/a (destructor)teedataobject_dealloc, /* tp_dealloc */
582n/a 0, /* tp_print */
583n/a 0, /* tp_getattr */
584n/a 0, /* tp_setattr */
585n/a 0, /* tp_reserved */
586n/a 0, /* tp_repr */
587n/a 0, /* tp_as_number */
588n/a 0, /* tp_as_sequence */
589n/a 0, /* tp_as_mapping */
590n/a 0, /* tp_hash */
591n/a 0, /* tp_call */
592n/a 0, /* tp_str */
593n/a PyObject_GenericGetAttr, /* tp_getattro */
594n/a 0, /* tp_setattro */
595n/a 0, /* tp_as_buffer */
596n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
597n/a teedataobject_doc, /* tp_doc */
598n/a (traverseproc)teedataobject_traverse, /* tp_traverse */
599n/a (inquiry)teedataobject_clear, /* tp_clear */
600n/a 0, /* tp_richcompare */
601n/a 0, /* tp_weaklistoffset */
602n/a 0, /* tp_iter */
603n/a 0, /* tp_iternext */
604n/a teedataobject_methods, /* tp_methods */
605n/a 0, /* tp_members */
606n/a 0, /* tp_getset */
607n/a 0, /* tp_base */
608n/a 0, /* tp_dict */
609n/a 0, /* tp_descr_get */
610n/a 0, /* tp_descr_set */
611n/a 0, /* tp_dictoffset */
612n/a 0, /* tp_init */
613n/a 0, /* tp_alloc */
614n/a teedataobject_new, /* tp_new */
615n/a PyObject_GC_Del, /* tp_free */
616n/a};
617n/a
618n/a
619n/astatic PyTypeObject tee_type;
620n/a
621n/astatic PyObject *
622n/atee_next(teeobject *to)
623n/a{
624n/a PyObject *value, *link;
625n/a
626n/a if (to->index >= LINKCELLS) {
627n/a link = teedataobject_jumplink(to->dataobj);
628n/a if (link == NULL)
629n/a return NULL;
630n/a Py_SETREF(to->dataobj, (teedataobject *)link);
631n/a to->index = 0;
632n/a }
633n/a value = teedataobject_getitem(to->dataobj, to->index);
634n/a if (value == NULL)
635n/a return NULL;
636n/a to->index++;
637n/a return value;
638n/a}
639n/a
640n/astatic int
641n/atee_traverse(teeobject *to, visitproc visit, void *arg)
642n/a{
643n/a Py_VISIT((PyObject *)to->dataobj);
644n/a return 0;
645n/a}
646n/a
647n/astatic PyObject *
648n/atee_copy(teeobject *to)
649n/a{
650n/a teeobject *newto;
651n/a
652n/a newto = PyObject_GC_New(teeobject, &tee_type);
653n/a if (newto == NULL)
654n/a return NULL;
655n/a Py_INCREF(to->dataobj);
656n/a newto->dataobj = to->dataobj;
657n/a newto->index = to->index;
658n/a newto->weakreflist = NULL;
659n/a PyObject_GC_Track(newto);
660n/a return (PyObject *)newto;
661n/a}
662n/a
663n/aPyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
664n/a
665n/astatic PyObject *
666n/atee_fromiterable(PyObject *iterable)
667n/a{
668n/a teeobject *to;
669n/a PyObject *it = NULL;
670n/a
671n/a it = PyObject_GetIter(iterable);
672n/a if (it == NULL)
673n/a return NULL;
674n/a if (PyObject_TypeCheck(it, &tee_type)) {
675n/a to = (teeobject *)tee_copy((teeobject *)it);
676n/a goto done;
677n/a }
678n/a
679n/a to = PyObject_GC_New(teeobject, &tee_type);
680n/a if (to == NULL)
681n/a goto done;
682n/a to->dataobj = (teedataobject *)teedataobject_newinternal(it);
683n/a if (!to->dataobj) {
684n/a PyObject_GC_Del(to);
685n/a to = NULL;
686n/a goto done;
687n/a }
688n/a
689n/a to->index = 0;
690n/a to->weakreflist = NULL;
691n/a PyObject_GC_Track(to);
692n/adone:
693n/a Py_XDECREF(it);
694n/a return (PyObject *)to;
695n/a}
696n/a
697n/astatic PyObject *
698n/atee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
699n/a{
700n/a PyObject *iterable;
701n/a
702n/a if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
703n/a return NULL;
704n/a return tee_fromiterable(iterable);
705n/a}
706n/a
707n/astatic int
708n/atee_clear(teeobject *to)
709n/a{
710n/a if (to->weakreflist != NULL)
711n/a PyObject_ClearWeakRefs((PyObject *) to);
712n/a Py_CLEAR(to->dataobj);
713n/a return 0;
714n/a}
715n/a
716n/astatic void
717n/atee_dealloc(teeobject *to)
718n/a{
719n/a PyObject_GC_UnTrack(to);
720n/a tee_clear(to);
721n/a PyObject_GC_Del(to);
722n/a}
723n/a
724n/astatic PyObject *
725n/atee_reduce(teeobject *to)
726n/a{
727n/a return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
728n/a}
729n/a
730n/astatic PyObject *
731n/atee_setstate(teeobject *to, PyObject *state)
732n/a{
733n/a teedataobject *tdo;
734n/a int index;
735n/a if (!PyTuple_Check(state)) {
736n/a PyErr_SetString(PyExc_TypeError, "state is not a tuple");
737n/a return NULL;
738n/a }
739n/a if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
740n/a return NULL;
741n/a }
742n/a if (index < 0 || index > LINKCELLS) {
743n/a PyErr_SetString(PyExc_ValueError, "Index out of range");
744n/a return NULL;
745n/a }
746n/a Py_INCREF(tdo);
747n/a Py_XSETREF(to->dataobj, tdo);
748n/a to->index = index;
749n/a Py_RETURN_NONE;
750n/a}
751n/a
752n/aPyDoc_STRVAR(teeobject_doc,
753n/a"Iterator wrapped to make it copyable");
754n/a
755n/astatic PyMethodDef tee_methods[] = {
756n/a {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
757n/a {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
758n/a {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
759n/a {NULL, NULL} /* sentinel */
760n/a};
761n/a
762n/astatic PyTypeObject tee_type = {
763n/a PyVarObject_HEAD_INIT(NULL, 0)
764n/a "itertools._tee", /* tp_name */
765n/a sizeof(teeobject), /* tp_basicsize */
766n/a 0, /* tp_itemsize */
767n/a /* methods */
768n/a (destructor)tee_dealloc, /* tp_dealloc */
769n/a 0, /* tp_print */
770n/a 0, /* tp_getattr */
771n/a 0, /* tp_setattr */
772n/a 0, /* tp_reserved */
773n/a 0, /* tp_repr */
774n/a 0, /* tp_as_number */
775n/a 0, /* tp_as_sequence */
776n/a 0, /* tp_as_mapping */
777n/a 0, /* tp_hash */
778n/a 0, /* tp_call */
779n/a 0, /* tp_str */
780n/a 0, /* tp_getattro */
781n/a 0, /* tp_setattro */
782n/a 0, /* tp_as_buffer */
783n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
784n/a teeobject_doc, /* tp_doc */
785n/a (traverseproc)tee_traverse, /* tp_traverse */
786n/a (inquiry)tee_clear, /* tp_clear */
787n/a 0, /* tp_richcompare */
788n/a offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
789n/a PyObject_SelfIter, /* tp_iter */
790n/a (iternextfunc)tee_next, /* tp_iternext */
791n/a tee_methods, /* tp_methods */
792n/a 0, /* tp_members */
793n/a 0, /* tp_getset */
794n/a 0, /* tp_base */
795n/a 0, /* tp_dict */
796n/a 0, /* tp_descr_get */
797n/a 0, /* tp_descr_set */
798n/a 0, /* tp_dictoffset */
799n/a 0, /* tp_init */
800n/a 0, /* tp_alloc */
801n/a tee_new, /* tp_new */
802n/a PyObject_GC_Del, /* tp_free */
803n/a};
804n/a
805n/astatic PyObject *
806n/atee(PyObject *self, PyObject *args)
807n/a{
808n/a Py_ssize_t i, n=2;
809n/a PyObject *it, *iterable, *copyable, *result;
810n/a _Py_IDENTIFIER(__copy__);
811n/a
812n/a if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
813n/a return NULL;
814n/a if (n < 0) {
815n/a PyErr_SetString(PyExc_ValueError, "n must be >= 0");
816n/a return NULL;
817n/a }
818n/a result = PyTuple_New(n);
819n/a if (result == NULL)
820n/a return NULL;
821n/a if (n == 0)
822n/a return result;
823n/a it = PyObject_GetIter(iterable);
824n/a if (it == NULL) {
825n/a Py_DECREF(result);
826n/a return NULL;
827n/a }
828n/a if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
829n/a copyable = tee_fromiterable(it);
830n/a Py_DECREF(it);
831n/a if (copyable == NULL) {
832n/a Py_DECREF(result);
833n/a return NULL;
834n/a }
835n/a } else
836n/a copyable = it;
837n/a PyTuple_SET_ITEM(result, 0, copyable);
838n/a for (i=1 ; i<n ; i++) {
839n/a
840n/a copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
841n/a if (copyable == NULL) {
842n/a Py_DECREF(result);
843n/a return NULL;
844n/a }
845n/a PyTuple_SET_ITEM(result, i, copyable);
846n/a }
847n/a return result;
848n/a}
849n/a
850n/aPyDoc_STRVAR(tee_doc,
851n/a"tee(iterable, n=2) --> tuple of n independent iterators.");
852n/a
853n/a
854n/a/* cycle object **************************************************************/
855n/a
856n/atypedef struct {
857n/a PyObject_HEAD
858n/a PyObject *it;
859n/a PyObject *saved;
860n/a Py_ssize_t index;
861n/a int firstpass;
862n/a} cycleobject;
863n/a
864n/astatic PyTypeObject cycle_type;
865n/a
866n/astatic PyObject *
867n/acycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
868n/a{
869n/a PyObject *it;
870n/a PyObject *iterable;
871n/a PyObject *saved;
872n/a cycleobject *lz;
873n/a
874n/a if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
875n/a return NULL;
876n/a
877n/a if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
878n/a return NULL;
879n/a
880n/a /* Get iterator. */
881n/a it = PyObject_GetIter(iterable);
882n/a if (it == NULL)
883n/a return NULL;
884n/a
885n/a saved = PyList_New(0);
886n/a if (saved == NULL) {
887n/a Py_DECREF(it);
888n/a return NULL;
889n/a }
890n/a
891n/a /* create cycleobject structure */
892n/a lz = (cycleobject *)type->tp_alloc(type, 0);
893n/a if (lz == NULL) {
894n/a Py_DECREF(it);
895n/a Py_DECREF(saved);
896n/a return NULL;
897n/a }
898n/a lz->it = it;
899n/a lz->saved = saved;
900n/a lz->index = 0;
901n/a lz->firstpass = 0;
902n/a
903n/a return (PyObject *)lz;
904n/a}
905n/a
906n/astatic void
907n/acycle_dealloc(cycleobject *lz)
908n/a{
909n/a PyObject_GC_UnTrack(lz);
910n/a Py_XDECREF(lz->it);
911n/a Py_XDECREF(lz->saved);
912n/a Py_TYPE(lz)->tp_free(lz);
913n/a}
914n/a
915n/astatic int
916n/acycle_traverse(cycleobject *lz, visitproc visit, void *arg)
917n/a{
918n/a if (lz->it)
919n/a Py_VISIT(lz->it);
920n/a Py_VISIT(lz->saved);
921n/a return 0;
922n/a}
923n/a
924n/astatic PyObject *
925n/acycle_next(cycleobject *lz)
926n/a{
927n/a PyObject *item;
928n/a
929n/a if (lz->it != NULL) {
930n/a item = PyIter_Next(lz->it);
931n/a if (item != NULL) {
932n/a if (lz->firstpass)
933n/a return item;
934n/a if (PyList_Append(lz->saved, item)) {
935n/a Py_DECREF(item);
936n/a return NULL;
937n/a }
938n/a return item;
939n/a }
940n/a /* Note: StopIteration is already cleared by PyIter_Next() */
941n/a if (PyErr_Occurred())
942n/a return NULL;
943n/a Py_CLEAR(lz->it);
944n/a }
945n/a if (Py_SIZE(lz->saved) == 0)
946n/a return NULL;
947n/a item = PyList_GET_ITEM(lz->saved, lz->index);
948n/a lz->index++;
949n/a if (lz->index >= Py_SIZE(lz->saved))
950n/a lz->index = 0;
951n/a Py_INCREF(item);
952n/a return item;
953n/a}
954n/a
955n/astatic PyObject *
956n/acycle_reduce(cycleobject *lz)
957n/a{
958n/a /* Create a new cycle with the iterator tuple, then set the saved state */
959n/a if (lz->it == NULL) {
960n/a PyObject *it = PyObject_GetIter(lz->saved);
961n/a if (it == NULL)
962n/a return NULL;
963n/a if (lz->index != 0) {
964n/a _Py_IDENTIFIER(__setstate__);
965n/a PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
966n/a "n", lz->index);
967n/a if (res == NULL) {
968n/a Py_DECREF(it);
969n/a return NULL;
970n/a }
971n/a Py_DECREF(res);
972n/a }
973n/a return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
974n/a }
975n/a return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
976n/a lz->firstpass);
977n/a}
978n/a
979n/astatic PyObject *
980n/acycle_setstate(cycleobject *lz, PyObject *state)
981n/a{
982n/a PyObject *saved=NULL;
983n/a int firstpass;
984n/a if (!PyTuple_Check(state)) {
985n/a PyErr_SetString(PyExc_TypeError, "state is not a tuple");
986n/a return NULL;
987n/a }
988n/a if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
989n/a return NULL;
990n/a }
991n/a Py_INCREF(saved);
992n/a Py_XSETREF(lz->saved, saved);
993n/a lz->firstpass = firstpass != 0;
994n/a lz->index = 0;
995n/a Py_RETURN_NONE;
996n/a}
997n/a
998n/astatic PyMethodDef cycle_methods[] = {
999n/a {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1000n/a reduce_doc},
1001n/a {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1002n/a setstate_doc},
1003n/a {NULL, NULL} /* sentinel */
1004n/a};
1005n/a
1006n/aPyDoc_STRVAR(cycle_doc,
1007n/a"cycle(iterable) --> cycle object\n\
1008n/a\n\
1009n/aReturn elements from the iterable until it is exhausted.\n\
1010n/aThen repeat the sequence indefinitely.");
1011n/a
1012n/astatic PyTypeObject cycle_type = {
1013n/a PyVarObject_HEAD_INIT(NULL, 0)
1014n/a "itertools.cycle", /* tp_name */
1015n/a sizeof(cycleobject), /* tp_basicsize */
1016n/a 0, /* tp_itemsize */
1017n/a /* methods */
1018n/a (destructor)cycle_dealloc, /* tp_dealloc */
1019n/a 0, /* tp_print */
1020n/a 0, /* tp_getattr */
1021n/a 0, /* tp_setattr */
1022n/a 0, /* tp_reserved */
1023n/a 0, /* tp_repr */
1024n/a 0, /* tp_as_number */
1025n/a 0, /* tp_as_sequence */
1026n/a 0, /* tp_as_mapping */
1027n/a 0, /* tp_hash */
1028n/a 0, /* tp_call */
1029n/a 0, /* tp_str */
1030n/a PyObject_GenericGetAttr, /* tp_getattro */
1031n/a 0, /* tp_setattro */
1032n/a 0, /* tp_as_buffer */
1033n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1034n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
1035n/a cycle_doc, /* tp_doc */
1036n/a (traverseproc)cycle_traverse, /* tp_traverse */
1037n/a 0, /* tp_clear */
1038n/a 0, /* tp_richcompare */
1039n/a 0, /* tp_weaklistoffset */
1040n/a PyObject_SelfIter, /* tp_iter */
1041n/a (iternextfunc)cycle_next, /* tp_iternext */
1042n/a cycle_methods, /* tp_methods */
1043n/a 0, /* tp_members */
1044n/a 0, /* tp_getset */
1045n/a 0, /* tp_base */
1046n/a 0, /* tp_dict */
1047n/a 0, /* tp_descr_get */
1048n/a 0, /* tp_descr_set */
1049n/a 0, /* tp_dictoffset */
1050n/a 0, /* tp_init */
1051n/a 0, /* tp_alloc */
1052n/a cycle_new, /* tp_new */
1053n/a PyObject_GC_Del, /* tp_free */
1054n/a};
1055n/a
1056n/a
1057n/a/* dropwhile object **********************************************************/
1058n/a
1059n/atypedef struct {
1060n/a PyObject_HEAD
1061n/a PyObject *func;
1062n/a PyObject *it;
1063n/a long start;
1064n/a} dropwhileobject;
1065n/a
1066n/astatic PyTypeObject dropwhile_type;
1067n/a
1068n/astatic PyObject *
1069n/adropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1070n/a{
1071n/a PyObject *func, *seq;
1072n/a PyObject *it;
1073n/a dropwhileobject *lz;
1074n/a
1075n/a if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1076n/a return NULL;
1077n/a
1078n/a if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1079n/a return NULL;
1080n/a
1081n/a /* Get iterator. */
1082n/a it = PyObject_GetIter(seq);
1083n/a if (it == NULL)
1084n/a return NULL;
1085n/a
1086n/a /* create dropwhileobject structure */
1087n/a lz = (dropwhileobject *)type->tp_alloc(type, 0);
1088n/a if (lz == NULL) {
1089n/a Py_DECREF(it);
1090n/a return NULL;
1091n/a }
1092n/a Py_INCREF(func);
1093n/a lz->func = func;
1094n/a lz->it = it;
1095n/a lz->start = 0;
1096n/a
1097n/a return (PyObject *)lz;
1098n/a}
1099n/a
1100n/astatic void
1101n/adropwhile_dealloc(dropwhileobject *lz)
1102n/a{
1103n/a PyObject_GC_UnTrack(lz);
1104n/a Py_XDECREF(lz->func);
1105n/a Py_XDECREF(lz->it);
1106n/a Py_TYPE(lz)->tp_free(lz);
1107n/a}
1108n/a
1109n/astatic int
1110n/adropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1111n/a{
1112n/a Py_VISIT(lz->it);
1113n/a Py_VISIT(lz->func);
1114n/a return 0;
1115n/a}
1116n/a
1117n/astatic PyObject *
1118n/adropwhile_next(dropwhileobject *lz)
1119n/a{
1120n/a PyObject *item, *good;
1121n/a PyObject *it = lz->it;
1122n/a long ok;
1123n/a PyObject *(*iternext)(PyObject *);
1124n/a
1125n/a iternext = *Py_TYPE(it)->tp_iternext;
1126n/a for (;;) {
1127n/a item = iternext(it);
1128n/a if (item == NULL)
1129n/a return NULL;
1130n/a if (lz->start == 1)
1131n/a return item;
1132n/a
1133n/a good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1134n/a if (good == NULL) {
1135n/a Py_DECREF(item);
1136n/a return NULL;
1137n/a }
1138n/a ok = PyObject_IsTrue(good);
1139n/a Py_DECREF(good);
1140n/a if (ok == 0) {
1141n/a lz->start = 1;
1142n/a return item;
1143n/a }
1144n/a Py_DECREF(item);
1145n/a if (ok < 0)
1146n/a return NULL;
1147n/a }
1148n/a}
1149n/a
1150n/astatic PyObject *
1151n/adropwhile_reduce(dropwhileobject *lz)
1152n/a{
1153n/a return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
1154n/a}
1155n/a
1156n/astatic PyObject *
1157n/adropwhile_setstate(dropwhileobject *lz, PyObject *state)
1158n/a{
1159n/a int start = PyObject_IsTrue(state);
1160n/a if (start < 0)
1161n/a return NULL;
1162n/a lz->start = start;
1163n/a Py_RETURN_NONE;
1164n/a}
1165n/a
1166n/astatic PyMethodDef dropwhile_methods[] = {
1167n/a {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1168n/a reduce_doc},
1169n/a {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1170n/a setstate_doc},
1171n/a {NULL, NULL} /* sentinel */
1172n/a};
1173n/a
1174n/aPyDoc_STRVAR(dropwhile_doc,
1175n/a"dropwhile(predicate, iterable) --> dropwhile object\n\
1176n/a\n\
1177n/aDrop items from the iterable while predicate(item) is true.\n\
1178n/aAfterwards, return every element until the iterable is exhausted.");
1179n/a
1180n/astatic PyTypeObject dropwhile_type = {
1181n/a PyVarObject_HEAD_INIT(NULL, 0)
1182n/a "itertools.dropwhile", /* tp_name */
1183n/a sizeof(dropwhileobject), /* tp_basicsize */
1184n/a 0, /* tp_itemsize */
1185n/a /* methods */
1186n/a (destructor)dropwhile_dealloc, /* tp_dealloc */
1187n/a 0, /* tp_print */
1188n/a 0, /* tp_getattr */
1189n/a 0, /* tp_setattr */
1190n/a 0, /* tp_reserved */
1191n/a 0, /* tp_repr */
1192n/a 0, /* tp_as_number */
1193n/a 0, /* tp_as_sequence */
1194n/a 0, /* tp_as_mapping */
1195n/a 0, /* tp_hash */
1196n/a 0, /* tp_call */
1197n/a 0, /* tp_str */
1198n/a PyObject_GenericGetAttr, /* tp_getattro */
1199n/a 0, /* tp_setattro */
1200n/a 0, /* tp_as_buffer */
1201n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1202n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
1203n/a dropwhile_doc, /* tp_doc */
1204n/a (traverseproc)dropwhile_traverse, /* tp_traverse */
1205n/a 0, /* tp_clear */
1206n/a 0, /* tp_richcompare */
1207n/a 0, /* tp_weaklistoffset */
1208n/a PyObject_SelfIter, /* tp_iter */
1209n/a (iternextfunc)dropwhile_next, /* tp_iternext */
1210n/a dropwhile_methods, /* tp_methods */
1211n/a 0, /* tp_members */
1212n/a 0, /* tp_getset */
1213n/a 0, /* tp_base */
1214n/a 0, /* tp_dict */
1215n/a 0, /* tp_descr_get */
1216n/a 0, /* tp_descr_set */
1217n/a 0, /* tp_dictoffset */
1218n/a 0, /* tp_init */
1219n/a 0, /* tp_alloc */
1220n/a dropwhile_new, /* tp_new */
1221n/a PyObject_GC_Del, /* tp_free */
1222n/a};
1223n/a
1224n/a
1225n/a/* takewhile object **********************************************************/
1226n/a
1227n/atypedef struct {
1228n/a PyObject_HEAD
1229n/a PyObject *func;
1230n/a PyObject *it;
1231n/a long stop;
1232n/a} takewhileobject;
1233n/a
1234n/astatic PyTypeObject takewhile_type;
1235n/a
1236n/astatic PyObject *
1237n/atakewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1238n/a{
1239n/a PyObject *func, *seq;
1240n/a PyObject *it;
1241n/a takewhileobject *lz;
1242n/a
1243n/a if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1244n/a return NULL;
1245n/a
1246n/a if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1247n/a return NULL;
1248n/a
1249n/a /* Get iterator. */
1250n/a it = PyObject_GetIter(seq);
1251n/a if (it == NULL)
1252n/a return NULL;
1253n/a
1254n/a /* create takewhileobject structure */
1255n/a lz = (takewhileobject *)type->tp_alloc(type, 0);
1256n/a if (lz == NULL) {
1257n/a Py_DECREF(it);
1258n/a return NULL;
1259n/a }
1260n/a Py_INCREF(func);
1261n/a lz->func = func;
1262n/a lz->it = it;
1263n/a lz->stop = 0;
1264n/a
1265n/a return (PyObject *)lz;
1266n/a}
1267n/a
1268n/astatic void
1269n/atakewhile_dealloc(takewhileobject *lz)
1270n/a{
1271n/a PyObject_GC_UnTrack(lz);
1272n/a Py_XDECREF(lz->func);
1273n/a Py_XDECREF(lz->it);
1274n/a Py_TYPE(lz)->tp_free(lz);
1275n/a}
1276n/a
1277n/astatic int
1278n/atakewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1279n/a{
1280n/a Py_VISIT(lz->it);
1281n/a Py_VISIT(lz->func);
1282n/a return 0;
1283n/a}
1284n/a
1285n/astatic PyObject *
1286n/atakewhile_next(takewhileobject *lz)
1287n/a{
1288n/a PyObject *item, *good;
1289n/a PyObject *it = lz->it;
1290n/a long ok;
1291n/a
1292n/a if (lz->stop == 1)
1293n/a return NULL;
1294n/a
1295n/a item = (*Py_TYPE(it)->tp_iternext)(it);
1296n/a if (item == NULL)
1297n/a return NULL;
1298n/a
1299n/a good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1300n/a if (good == NULL) {
1301n/a Py_DECREF(item);
1302n/a return NULL;
1303n/a }
1304n/a ok = PyObject_IsTrue(good);
1305n/a Py_DECREF(good);
1306n/a if (ok > 0)
1307n/a return item;
1308n/a Py_DECREF(item);
1309n/a if (ok == 0)
1310n/a lz->stop = 1;
1311n/a return NULL;
1312n/a}
1313n/a
1314n/astatic PyObject *
1315n/atakewhile_reduce(takewhileobject *lz)
1316n/a{
1317n/a return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
1318n/a}
1319n/a
1320n/astatic PyObject *
1321n/atakewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1322n/a{
1323n/a int stop = PyObject_IsTrue(state);
1324n/a
1325n/a if (stop < 0)
1326n/a return NULL;
1327n/a lz->stop = stop;
1328n/a Py_RETURN_NONE;
1329n/a}
1330n/a
1331n/astatic PyMethodDef takewhile_reduce_methods[] = {
1332n/a {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1333n/a reduce_doc},
1334n/a {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1335n/a setstate_doc},
1336n/a {NULL, NULL} /* sentinel */
1337n/a};
1338n/aPyDoc_STRVAR(takewhile_doc,
1339n/a"takewhile(predicate, iterable) --> takewhile object\n\
1340n/a\n\
1341n/aReturn successive entries from an iterable as long as the \n\
1342n/apredicate evaluates to true for each entry.");
1343n/a
1344n/astatic PyTypeObject takewhile_type = {
1345n/a PyVarObject_HEAD_INIT(NULL, 0)
1346n/a "itertools.takewhile", /* tp_name */
1347n/a sizeof(takewhileobject), /* tp_basicsize */
1348n/a 0, /* tp_itemsize */
1349n/a /* methods */
1350n/a (destructor)takewhile_dealloc, /* tp_dealloc */
1351n/a 0, /* tp_print */
1352n/a 0, /* tp_getattr */
1353n/a 0, /* tp_setattr */
1354n/a 0, /* tp_reserved */
1355n/a 0, /* tp_repr */
1356n/a 0, /* tp_as_number */
1357n/a 0, /* tp_as_sequence */
1358n/a 0, /* tp_as_mapping */
1359n/a 0, /* tp_hash */
1360n/a 0, /* tp_call */
1361n/a 0, /* tp_str */
1362n/a PyObject_GenericGetAttr, /* tp_getattro */
1363n/a 0, /* tp_setattro */
1364n/a 0, /* tp_as_buffer */
1365n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1366n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
1367n/a takewhile_doc, /* tp_doc */
1368n/a (traverseproc)takewhile_traverse, /* tp_traverse */
1369n/a 0, /* tp_clear */
1370n/a 0, /* tp_richcompare */
1371n/a 0, /* tp_weaklistoffset */
1372n/a PyObject_SelfIter, /* tp_iter */
1373n/a (iternextfunc)takewhile_next, /* tp_iternext */
1374n/a takewhile_reduce_methods, /* tp_methods */
1375n/a 0, /* tp_members */
1376n/a 0, /* tp_getset */
1377n/a 0, /* tp_base */
1378n/a 0, /* tp_dict */
1379n/a 0, /* tp_descr_get */
1380n/a 0, /* tp_descr_set */
1381n/a 0, /* tp_dictoffset */
1382n/a 0, /* tp_init */
1383n/a 0, /* tp_alloc */
1384n/a takewhile_new, /* tp_new */
1385n/a PyObject_GC_Del, /* tp_free */
1386n/a};
1387n/a
1388n/a
1389n/a/* islice object *************************************************************/
1390n/a
1391n/atypedef struct {
1392n/a PyObject_HEAD
1393n/a PyObject *it;
1394n/a Py_ssize_t next;
1395n/a Py_ssize_t stop;
1396n/a Py_ssize_t step;
1397n/a Py_ssize_t cnt;
1398n/a} isliceobject;
1399n/a
1400n/astatic PyTypeObject islice_type;
1401n/a
1402n/astatic PyObject *
1403n/aislice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1404n/a{
1405n/a PyObject *seq;
1406n/a Py_ssize_t start=0, stop=-1, step=1;
1407n/a PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1408n/a Py_ssize_t numargs;
1409n/a isliceobject *lz;
1410n/a
1411n/a if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1412n/a return NULL;
1413n/a
1414n/a if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1415n/a return NULL;
1416n/a
1417n/a numargs = PyTuple_Size(args);
1418n/a if (numargs == 2) {
1419n/a if (a1 != Py_None) {
1420n/a stop = PyLong_AsSsize_t(a1);
1421n/a if (stop == -1) {
1422n/a if (PyErr_Occurred())
1423n/a PyErr_Clear();
1424n/a PyErr_SetString(PyExc_ValueError,
1425n/a "Stop argument for islice() must be None or "
1426n/a "an integer: 0 <= x <= sys.maxsize.");
1427n/a return NULL;
1428n/a }
1429n/a }
1430n/a } else {
1431n/a if (a1 != Py_None)
1432n/a start = PyLong_AsSsize_t(a1);
1433n/a if (start == -1 && PyErr_Occurred())
1434n/a PyErr_Clear();
1435n/a if (a2 != Py_None) {
1436n/a stop = PyLong_AsSsize_t(a2);
1437n/a if (stop == -1) {
1438n/a if (PyErr_Occurred())
1439n/a PyErr_Clear();
1440n/a PyErr_SetString(PyExc_ValueError,
1441n/a "Stop argument for islice() must be None or "
1442n/a "an integer: 0 <= x <= sys.maxsize.");
1443n/a return NULL;
1444n/a }
1445n/a }
1446n/a }
1447n/a if (start<0 || stop<-1) {
1448n/a PyErr_SetString(PyExc_ValueError,
1449n/a "Indices for islice() must be None or "
1450n/a "an integer: 0 <= x <= sys.maxsize.");
1451n/a return NULL;
1452n/a }
1453n/a
1454n/a if (a3 != NULL) {
1455n/a if (a3 != Py_None)
1456n/a step = PyLong_AsSsize_t(a3);
1457n/a if (step == -1 && PyErr_Occurred())
1458n/a PyErr_Clear();
1459n/a }
1460n/a if (step<1) {
1461n/a PyErr_SetString(PyExc_ValueError,
1462n/a "Step for islice() must be a positive integer or None.");
1463n/a return NULL;
1464n/a }
1465n/a
1466n/a /* Get iterator. */
1467n/a it = PyObject_GetIter(seq);
1468n/a if (it == NULL)
1469n/a return NULL;
1470n/a
1471n/a /* create isliceobject structure */
1472n/a lz = (isliceobject *)type->tp_alloc(type, 0);
1473n/a if (lz == NULL) {
1474n/a Py_DECREF(it);
1475n/a return NULL;
1476n/a }
1477n/a lz->it = it;
1478n/a lz->next = start;
1479n/a lz->stop = stop;
1480n/a lz->step = step;
1481n/a lz->cnt = 0L;
1482n/a
1483n/a return (PyObject *)lz;
1484n/a}
1485n/a
1486n/astatic void
1487n/aislice_dealloc(isliceobject *lz)
1488n/a{
1489n/a PyObject_GC_UnTrack(lz);
1490n/a Py_XDECREF(lz->it);
1491n/a Py_TYPE(lz)->tp_free(lz);
1492n/a}
1493n/a
1494n/astatic int
1495n/aislice_traverse(isliceobject *lz, visitproc visit, void *arg)
1496n/a{
1497n/a Py_VISIT(lz->it);
1498n/a return 0;
1499n/a}
1500n/a
1501n/astatic PyObject *
1502n/aislice_next(isliceobject *lz)
1503n/a{
1504n/a PyObject *item;
1505n/a PyObject *it = lz->it;
1506n/a Py_ssize_t stop = lz->stop;
1507n/a Py_ssize_t oldnext;
1508n/a PyObject *(*iternext)(PyObject *);
1509n/a
1510n/a if (it == NULL)
1511n/a return NULL;
1512n/a
1513n/a iternext = *Py_TYPE(it)->tp_iternext;
1514n/a while (lz->cnt < lz->next) {
1515n/a item = iternext(it);
1516n/a if (item == NULL)
1517n/a goto empty;
1518n/a Py_DECREF(item);
1519n/a lz->cnt++;
1520n/a }
1521n/a if (stop != -1 && lz->cnt >= stop)
1522n/a goto empty;
1523n/a item = iternext(it);
1524n/a if (item == NULL)
1525n/a goto empty;
1526n/a lz->cnt++;
1527n/a oldnext = lz->next;
1528n/a /* The (size_t) cast below avoids the danger of undefined
1529n/a behaviour from signed integer overflow. */
1530n/a lz->next += (size_t)lz->step;
1531n/a if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1532n/a lz->next = stop;
1533n/a return item;
1534n/a
1535n/aempty:
1536n/a Py_CLEAR(lz->it);
1537n/a return NULL;
1538n/a}
1539n/a
1540n/astatic PyObject *
1541n/aislice_reduce(isliceobject *lz)
1542n/a{
1543n/a /* When unpickled, generate a new object with the same bounds,
1544n/a * then 'setstate' with the next and count
1545n/a */
1546n/a PyObject *stop;
1547n/a
1548n/a if (lz->it == NULL) {
1549n/a PyObject *empty_list;
1550n/a PyObject *empty_it;
1551n/a empty_list = PyList_New(0);
1552n/a if (empty_list == NULL)
1553n/a return NULL;
1554n/a empty_it = PyObject_GetIter(empty_list);
1555n/a Py_DECREF(empty_list);
1556n/a if (empty_it == NULL)
1557n/a return NULL;
1558n/a return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1559n/a }
1560n/a if (lz->stop == -1) {
1561n/a stop = Py_None;
1562n/a Py_INCREF(stop);
1563n/a } else {
1564n/a stop = PyLong_FromSsize_t(lz->stop);
1565n/a if (stop == NULL)
1566n/a return NULL;
1567n/a }
1568n/a return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1569n/a lz->it, lz->next, stop, lz->step,
1570n/a lz->cnt);
1571n/a}
1572n/a
1573n/astatic PyObject *
1574n/aislice_setstate(isliceobject *lz, PyObject *state)
1575n/a{
1576n/a Py_ssize_t cnt = PyLong_AsSsize_t(state);
1577n/a
1578n/a if (cnt == -1 && PyErr_Occurred())
1579n/a return NULL;
1580n/a lz->cnt = cnt;
1581n/a Py_RETURN_NONE;
1582n/a}
1583n/a
1584n/astatic PyMethodDef islice_methods[] = {
1585n/a {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1586n/a reduce_doc},
1587n/a {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1588n/a setstate_doc},
1589n/a {NULL, NULL} /* sentinel */
1590n/a};
1591n/a
1592n/aPyDoc_STRVAR(islice_doc,
1593n/a"islice(iterable, stop) --> islice object\n\
1594n/aislice(iterable, start, stop[, step]) --> islice object\n\
1595n/a\n\
1596n/aReturn an iterator whose next() method returns selected values from an\n\
1597n/aiterable. If start is specified, will skip all preceding elements;\n\
1598n/aotherwise, start defaults to zero. Step defaults to one. If\n\
1599n/aspecified as another value, step determines how many values are \n\
1600n/askipped between successive calls. Works like a slice() on a list\n\
1601n/abut returns an iterator.");
1602n/a
1603n/astatic PyTypeObject islice_type = {
1604n/a PyVarObject_HEAD_INIT(NULL, 0)
1605n/a "itertools.islice", /* tp_name */
1606n/a sizeof(isliceobject), /* tp_basicsize */
1607n/a 0, /* tp_itemsize */
1608n/a /* methods */
1609n/a (destructor)islice_dealloc, /* tp_dealloc */
1610n/a 0, /* tp_print */
1611n/a 0, /* tp_getattr */
1612n/a 0, /* tp_setattr */
1613n/a 0, /* tp_reserved */
1614n/a 0, /* tp_repr */
1615n/a 0, /* tp_as_number */
1616n/a 0, /* tp_as_sequence */
1617n/a 0, /* tp_as_mapping */
1618n/a 0, /* tp_hash */
1619n/a 0, /* tp_call */
1620n/a 0, /* tp_str */
1621n/a PyObject_GenericGetAttr, /* tp_getattro */
1622n/a 0, /* tp_setattro */
1623n/a 0, /* tp_as_buffer */
1624n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1625n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
1626n/a islice_doc, /* tp_doc */
1627n/a (traverseproc)islice_traverse, /* tp_traverse */
1628n/a 0, /* tp_clear */
1629n/a 0, /* tp_richcompare */
1630n/a 0, /* tp_weaklistoffset */
1631n/a PyObject_SelfIter, /* tp_iter */
1632n/a (iternextfunc)islice_next, /* tp_iternext */
1633n/a islice_methods, /* tp_methods */
1634n/a 0, /* tp_members */
1635n/a 0, /* tp_getset */
1636n/a 0, /* tp_base */
1637n/a 0, /* tp_dict */
1638n/a 0, /* tp_descr_get */
1639n/a 0, /* tp_descr_set */
1640n/a 0, /* tp_dictoffset */
1641n/a 0, /* tp_init */
1642n/a 0, /* tp_alloc */
1643n/a islice_new, /* tp_new */
1644n/a PyObject_GC_Del, /* tp_free */
1645n/a};
1646n/a
1647n/a
1648n/a/* starmap object ************************************************************/
1649n/a
1650n/atypedef struct {
1651n/a PyObject_HEAD
1652n/a PyObject *func;
1653n/a PyObject *it;
1654n/a} starmapobject;
1655n/a
1656n/astatic PyTypeObject starmap_type;
1657n/a
1658n/astatic PyObject *
1659n/astarmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1660n/a{
1661n/a PyObject *func, *seq;
1662n/a PyObject *it;
1663n/a starmapobject *lz;
1664n/a
1665n/a if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1666n/a return NULL;
1667n/a
1668n/a if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1669n/a return NULL;
1670n/a
1671n/a /* Get iterator. */
1672n/a it = PyObject_GetIter(seq);
1673n/a if (it == NULL)
1674n/a return NULL;
1675n/a
1676n/a /* create starmapobject structure */
1677n/a lz = (starmapobject *)type->tp_alloc(type, 0);
1678n/a if (lz == NULL) {
1679n/a Py_DECREF(it);
1680n/a return NULL;
1681n/a }
1682n/a Py_INCREF(func);
1683n/a lz->func = func;
1684n/a lz->it = it;
1685n/a
1686n/a return (PyObject *)lz;
1687n/a}
1688n/a
1689n/astatic void
1690n/astarmap_dealloc(starmapobject *lz)
1691n/a{
1692n/a PyObject_GC_UnTrack(lz);
1693n/a Py_XDECREF(lz->func);
1694n/a Py_XDECREF(lz->it);
1695n/a Py_TYPE(lz)->tp_free(lz);
1696n/a}
1697n/a
1698n/astatic int
1699n/astarmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1700n/a{
1701n/a Py_VISIT(lz->it);
1702n/a Py_VISIT(lz->func);
1703n/a return 0;
1704n/a}
1705n/a
1706n/astatic PyObject *
1707n/astarmap_next(starmapobject *lz)
1708n/a{
1709n/a PyObject *args;
1710n/a PyObject *result;
1711n/a PyObject *it = lz->it;
1712n/a
1713n/a args = (*Py_TYPE(it)->tp_iternext)(it);
1714n/a if (args == NULL)
1715n/a return NULL;
1716n/a if (!PyTuple_CheckExact(args)) {
1717n/a PyObject *newargs = PySequence_Tuple(args);
1718n/a Py_DECREF(args);
1719n/a if (newargs == NULL)
1720n/a return NULL;
1721n/a args = newargs;
1722n/a }
1723n/a result = PyObject_Call(lz->func, args, NULL);
1724n/a Py_DECREF(args);
1725n/a return result;
1726n/a}
1727n/a
1728n/astatic PyObject *
1729n/astarmap_reduce(starmapobject *lz)
1730n/a{
1731n/a /* Just pickle the iterator */
1732n/a return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1733n/a}
1734n/a
1735n/astatic PyMethodDef starmap_methods[] = {
1736n/a {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1737n/a reduce_doc},
1738n/a {NULL, NULL} /* sentinel */
1739n/a};
1740n/a
1741n/aPyDoc_STRVAR(starmap_doc,
1742n/a"starmap(function, sequence) --> starmap object\n\
1743n/a\n\
1744n/aReturn an iterator whose values are returned from the function evaluated\n\
1745n/awith an argument tuple taken from the given sequence.");
1746n/a
1747n/astatic PyTypeObject starmap_type = {
1748n/a PyVarObject_HEAD_INIT(NULL, 0)
1749n/a "itertools.starmap", /* tp_name */
1750n/a sizeof(starmapobject), /* tp_basicsize */
1751n/a 0, /* tp_itemsize */
1752n/a /* methods */
1753n/a (destructor)starmap_dealloc, /* tp_dealloc */
1754n/a 0, /* tp_print */
1755n/a 0, /* tp_getattr */
1756n/a 0, /* tp_setattr */
1757n/a 0, /* tp_reserved */
1758n/a 0, /* tp_repr */
1759n/a 0, /* tp_as_number */
1760n/a 0, /* tp_as_sequence */
1761n/a 0, /* tp_as_mapping */
1762n/a 0, /* tp_hash */
1763n/a 0, /* tp_call */
1764n/a 0, /* tp_str */
1765n/a PyObject_GenericGetAttr, /* tp_getattro */
1766n/a 0, /* tp_setattro */
1767n/a 0, /* tp_as_buffer */
1768n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1769n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
1770n/a starmap_doc, /* tp_doc */
1771n/a (traverseproc)starmap_traverse, /* tp_traverse */
1772n/a 0, /* tp_clear */
1773n/a 0, /* tp_richcompare */
1774n/a 0, /* tp_weaklistoffset */
1775n/a PyObject_SelfIter, /* tp_iter */
1776n/a (iternextfunc)starmap_next, /* tp_iternext */
1777n/a starmap_methods, /* tp_methods */
1778n/a 0, /* tp_members */
1779n/a 0, /* tp_getset */
1780n/a 0, /* tp_base */
1781n/a 0, /* tp_dict */
1782n/a 0, /* tp_descr_get */
1783n/a 0, /* tp_descr_set */
1784n/a 0, /* tp_dictoffset */
1785n/a 0, /* tp_init */
1786n/a 0, /* tp_alloc */
1787n/a starmap_new, /* tp_new */
1788n/a PyObject_GC_Del, /* tp_free */
1789n/a};
1790n/a
1791n/a
1792n/a/* chain object **************************************************************/
1793n/a
1794n/atypedef struct {
1795n/a PyObject_HEAD
1796n/a PyObject *source; /* Iterator over input iterables */
1797n/a PyObject *active; /* Currently running input iterator */
1798n/a} chainobject;
1799n/a
1800n/astatic PyTypeObject chain_type;
1801n/a
1802n/astatic PyObject *
1803n/achain_new_internal(PyTypeObject *type, PyObject *source)
1804n/a{
1805n/a chainobject *lz;
1806n/a
1807n/a lz = (chainobject *)type->tp_alloc(type, 0);
1808n/a if (lz == NULL) {
1809n/a Py_DECREF(source);
1810n/a return NULL;
1811n/a }
1812n/a
1813n/a lz->source = source;
1814n/a lz->active = NULL;
1815n/a return (PyObject *)lz;
1816n/a}
1817n/a
1818n/astatic PyObject *
1819n/achain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1820n/a{
1821n/a PyObject *source;
1822n/a
1823n/a if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1824n/a return NULL;
1825n/a
1826n/a source = PyObject_GetIter(args);
1827n/a if (source == NULL)
1828n/a return NULL;
1829n/a
1830n/a return chain_new_internal(type, source);
1831n/a}
1832n/a
1833n/astatic PyObject *
1834n/achain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1835n/a{
1836n/a PyObject *source;
1837n/a
1838n/a source = PyObject_GetIter(arg);
1839n/a if (source == NULL)
1840n/a return NULL;
1841n/a
1842n/a return chain_new_internal(type, source);
1843n/a}
1844n/a
1845n/astatic void
1846n/achain_dealloc(chainobject *lz)
1847n/a{
1848n/a PyObject_GC_UnTrack(lz);
1849n/a Py_XDECREF(lz->active);
1850n/a Py_XDECREF(lz->source);
1851n/a Py_TYPE(lz)->tp_free(lz);
1852n/a}
1853n/a
1854n/astatic int
1855n/achain_traverse(chainobject *lz, visitproc visit, void *arg)
1856n/a{
1857n/a Py_VISIT(lz->source);
1858n/a Py_VISIT(lz->active);
1859n/a return 0;
1860n/a}
1861n/a
1862n/astatic PyObject *
1863n/achain_next(chainobject *lz)
1864n/a{
1865n/a PyObject *item;
1866n/a
1867n/a if (lz->source == NULL)
1868n/a return NULL; /* already stopped */
1869n/a
1870n/a if (lz->active == NULL) {
1871n/a PyObject *iterable = PyIter_Next(lz->source);
1872n/a if (iterable == NULL) {
1873n/a Py_CLEAR(lz->source);
1874n/a return NULL; /* no more input sources */
1875n/a }
1876n/a lz->active = PyObject_GetIter(iterable);
1877n/a Py_DECREF(iterable);
1878n/a if (lz->active == NULL) {
1879n/a Py_CLEAR(lz->source);
1880n/a return NULL; /* input not iterable */
1881n/a }
1882n/a }
1883n/a item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1884n/a if (item != NULL)
1885n/a return item;
1886n/a if (PyErr_Occurred()) {
1887n/a if (PyErr_ExceptionMatches(PyExc_StopIteration))
1888n/a PyErr_Clear();
1889n/a else
1890n/a return NULL; /* input raised an exception */
1891n/a }
1892n/a Py_CLEAR(lz->active);
1893n/a return chain_next(lz); /* recurse and use next active */
1894n/a}
1895n/a
1896n/astatic PyObject *
1897n/achain_reduce(chainobject *lz)
1898n/a{
1899n/a if (lz->source) {
1900n/a /* we can't pickle function objects (itertools.from_iterable) so
1901n/a * we must use setstate to replace the iterable. One day we
1902n/a * will fix pickling of functions
1903n/a */
1904n/a if (lz->active) {
1905n/a return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1906n/a } else {
1907n/a return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1908n/a }
1909n/a } else {
1910n/a return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1911n/a }
1912n/a return NULL;
1913n/a}
1914n/a
1915n/astatic PyObject *
1916n/achain_setstate(chainobject *lz, PyObject *state)
1917n/a{
1918n/a PyObject *source, *active=NULL;
1919n/a
1920n/a if (!PyTuple_Check(state)) {
1921n/a PyErr_SetString(PyExc_TypeError, "state is not a tuple");
1922n/a return NULL;
1923n/a }
1924n/a if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
1925n/a return NULL;
1926n/a }
1927n/a if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
1928n/a PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
1929n/a return NULL;
1930n/a }
1931n/a
1932n/a Py_INCREF(source);
1933n/a Py_XSETREF(lz->source, source);
1934n/a Py_XINCREF(active);
1935n/a Py_XSETREF(lz->active, active);
1936n/a Py_RETURN_NONE;
1937n/a}
1938n/a
1939n/aPyDoc_STRVAR(chain_doc,
1940n/a"chain(*iterables) --> chain object\n\
1941n/a\n\
1942n/aReturn a chain object whose .__next__() method returns elements from the\n\
1943n/afirst iterable until it is exhausted, then elements from the next\n\
1944n/aiterable, until all of the iterables are exhausted.");
1945n/a
1946n/aPyDoc_STRVAR(chain_from_iterable_doc,
1947n/a"chain.from_iterable(iterable) --> chain object\n\
1948n/a\n\
1949n/aAlternate chain() constructor taking a single iterable argument\n\
1950n/athat evaluates lazily.");
1951n/a
1952n/astatic PyMethodDef chain_methods[] = {
1953n/a {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1954n/a chain_from_iterable_doc},
1955n/a {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1956n/a reduce_doc},
1957n/a {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1958n/a setstate_doc},
1959n/a {NULL, NULL} /* sentinel */
1960n/a};
1961n/a
1962n/astatic PyTypeObject chain_type = {
1963n/a PyVarObject_HEAD_INIT(NULL, 0)
1964n/a "itertools.chain", /* tp_name */
1965n/a sizeof(chainobject), /* tp_basicsize */
1966n/a 0, /* tp_itemsize */
1967n/a /* methods */
1968n/a (destructor)chain_dealloc, /* tp_dealloc */
1969n/a 0, /* tp_print */
1970n/a 0, /* tp_getattr */
1971n/a 0, /* tp_setattr */
1972n/a 0, /* tp_reserved */
1973n/a 0, /* tp_repr */
1974n/a 0, /* tp_as_number */
1975n/a 0, /* tp_as_sequence */
1976n/a 0, /* tp_as_mapping */
1977n/a 0, /* tp_hash */
1978n/a 0, /* tp_call */
1979n/a 0, /* tp_str */
1980n/a PyObject_GenericGetAttr, /* tp_getattro */
1981n/a 0, /* tp_setattro */
1982n/a 0, /* tp_as_buffer */
1983n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1984n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
1985n/a chain_doc, /* tp_doc */
1986n/a (traverseproc)chain_traverse, /* tp_traverse */
1987n/a 0, /* tp_clear */
1988n/a 0, /* tp_richcompare */
1989n/a 0, /* tp_weaklistoffset */
1990n/a PyObject_SelfIter, /* tp_iter */
1991n/a (iternextfunc)chain_next, /* tp_iternext */
1992n/a chain_methods, /* tp_methods */
1993n/a 0, /* tp_members */
1994n/a 0, /* tp_getset */
1995n/a 0, /* tp_base */
1996n/a 0, /* tp_dict */
1997n/a 0, /* tp_descr_get */
1998n/a 0, /* tp_descr_set */
1999n/a 0, /* tp_dictoffset */
2000n/a 0, /* tp_init */
2001n/a 0, /* tp_alloc */
2002n/a chain_new, /* tp_new */
2003n/a PyObject_GC_Del, /* tp_free */
2004n/a};
2005n/a
2006n/a
2007n/a/* product object ************************************************************/
2008n/a
2009n/atypedef struct {
2010n/a PyObject_HEAD
2011n/a PyObject *pools; /* tuple of pool tuples */
2012n/a Py_ssize_t *indices; /* one index per pool */
2013n/a PyObject *result; /* most recently returned result tuple */
2014n/a int stopped; /* set to 1 when the iterator is exhausted */
2015n/a} productobject;
2016n/a
2017n/astatic PyTypeObject product_type;
2018n/a
2019n/astatic PyObject *
2020n/aproduct_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2021n/a{
2022n/a productobject *lz;
2023n/a Py_ssize_t nargs, npools, repeat=1;
2024n/a PyObject *pools = NULL;
2025n/a Py_ssize_t *indices = NULL;
2026n/a Py_ssize_t i;
2027n/a
2028n/a if (kwds != NULL) {
2029n/a char *kwlist[] = {"repeat", 0};
2030n/a PyObject *tmpargs = PyTuple_New(0);
2031n/a if (tmpargs == NULL)
2032n/a return NULL;
2033n/a if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2034n/a kwlist, &repeat)) {
2035n/a Py_DECREF(tmpargs);
2036n/a return NULL;
2037n/a }
2038n/a Py_DECREF(tmpargs);
2039n/a if (repeat < 0) {
2040n/a PyErr_SetString(PyExc_ValueError,
2041n/a "repeat argument cannot be negative");
2042n/a return NULL;
2043n/a }
2044n/a }
2045n/a
2046n/a assert(PyTuple_CheckExact(args));
2047n/a if (repeat == 0) {
2048n/a nargs = 0;
2049n/a } else {
2050n/a nargs = PyTuple_GET_SIZE(args);
2051n/a if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
2052n/a PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2053n/a return NULL;
2054n/a }
2055n/a }
2056n/a npools = nargs * repeat;
2057n/a
2058n/a indices = PyMem_New(Py_ssize_t, npools);
2059n/a if (indices == NULL) {
2060n/a PyErr_NoMemory();
2061n/a goto error;
2062n/a }
2063n/a
2064n/a pools = PyTuple_New(npools);
2065n/a if (pools == NULL)
2066n/a goto error;
2067n/a
2068n/a for (i=0; i < nargs ; ++i) {
2069n/a PyObject *item = PyTuple_GET_ITEM(args, i);
2070n/a PyObject *pool = PySequence_Tuple(item);
2071n/a if (pool == NULL)
2072n/a goto error;
2073n/a PyTuple_SET_ITEM(pools, i, pool);
2074n/a indices[i] = 0;
2075n/a }
2076n/a for ( ; i < npools; ++i) {
2077n/a PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2078n/a Py_INCREF(pool);
2079n/a PyTuple_SET_ITEM(pools, i, pool);
2080n/a indices[i] = 0;
2081n/a }
2082n/a
2083n/a /* create productobject structure */
2084n/a lz = (productobject *)type->tp_alloc(type, 0);
2085n/a if (lz == NULL)
2086n/a goto error;
2087n/a
2088n/a lz->pools = pools;
2089n/a lz->indices = indices;
2090n/a lz->result = NULL;
2091n/a lz->stopped = 0;
2092n/a
2093n/a return (PyObject *)lz;
2094n/a
2095n/aerror:
2096n/a if (indices != NULL)
2097n/a PyMem_Free(indices);
2098n/a Py_XDECREF(pools);
2099n/a return NULL;
2100n/a}
2101n/a
2102n/astatic void
2103n/aproduct_dealloc(productobject *lz)
2104n/a{
2105n/a PyObject_GC_UnTrack(lz);
2106n/a Py_XDECREF(lz->pools);
2107n/a Py_XDECREF(lz->result);
2108n/a if (lz->indices != NULL)
2109n/a PyMem_Free(lz->indices);
2110n/a Py_TYPE(lz)->tp_free(lz);
2111n/a}
2112n/a
2113n/astatic PyObject *
2114n/aproduct_sizeof(productobject *lz, void *unused)
2115n/a{
2116n/a Py_ssize_t res;
2117n/a
2118n/a res = _PyObject_SIZE(Py_TYPE(lz));
2119n/a res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2120n/a return PyLong_FromSsize_t(res);
2121n/a}
2122n/a
2123n/aPyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2124n/a
2125n/astatic int
2126n/aproduct_traverse(productobject *lz, visitproc visit, void *arg)
2127n/a{
2128n/a Py_VISIT(lz->pools);
2129n/a Py_VISIT(lz->result);
2130n/a return 0;
2131n/a}
2132n/a
2133n/astatic PyObject *
2134n/aproduct_next(productobject *lz)
2135n/a{
2136n/a PyObject *pool;
2137n/a PyObject *elem;
2138n/a PyObject *oldelem;
2139n/a PyObject *pools = lz->pools;
2140n/a PyObject *result = lz->result;
2141n/a Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2142n/a Py_ssize_t i;
2143n/a
2144n/a if (lz->stopped)
2145n/a return NULL;
2146n/a
2147n/a if (result == NULL) {
2148n/a /* On the first pass, return an initial tuple filled with the
2149n/a first element from each pool. */
2150n/a result = PyTuple_New(npools);
2151n/a if (result == NULL)
2152n/a goto empty;
2153n/a lz->result = result;
2154n/a for (i=0; i < npools; i++) {
2155n/a pool = PyTuple_GET_ITEM(pools, i);
2156n/a if (PyTuple_GET_SIZE(pool) == 0)
2157n/a goto empty;
2158n/a elem = PyTuple_GET_ITEM(pool, 0);
2159n/a Py_INCREF(elem);
2160n/a PyTuple_SET_ITEM(result, i, elem);
2161n/a }
2162n/a } else {
2163n/a Py_ssize_t *indices = lz->indices;
2164n/a
2165n/a /* Copy the previous result tuple or re-use it if available */
2166n/a if (Py_REFCNT(result) > 1) {
2167n/a PyObject *old_result = result;
2168n/a result = PyTuple_New(npools);
2169n/a if (result == NULL)
2170n/a goto empty;
2171n/a lz->result = result;
2172n/a for (i=0; i < npools; i++) {
2173n/a elem = PyTuple_GET_ITEM(old_result, i);
2174n/a Py_INCREF(elem);
2175n/a PyTuple_SET_ITEM(result, i, elem);
2176n/a }
2177n/a Py_DECREF(old_result);
2178n/a }
2179n/a /* Now, we've got the only copy so we can update it in-place */
2180n/a assert (npools==0 || Py_REFCNT(result) == 1);
2181n/a
2182n/a /* Update the pool indices right-to-left. Only advance to the
2183n/a next pool when the previous one rolls-over */
2184n/a for (i=npools-1 ; i >= 0 ; i--) {
2185n/a pool = PyTuple_GET_ITEM(pools, i);
2186n/a indices[i]++;
2187n/a if (indices[i] == PyTuple_GET_SIZE(pool)) {
2188n/a /* Roll-over and advance to next pool */
2189n/a indices[i] = 0;
2190n/a elem = PyTuple_GET_ITEM(pool, 0);
2191n/a Py_INCREF(elem);
2192n/a oldelem = PyTuple_GET_ITEM(result, i);
2193n/a PyTuple_SET_ITEM(result, i, elem);
2194n/a Py_DECREF(oldelem);
2195n/a } else {
2196n/a /* No rollover. Just increment and stop here. */
2197n/a elem = PyTuple_GET_ITEM(pool, indices[i]);
2198n/a Py_INCREF(elem);
2199n/a oldelem = PyTuple_GET_ITEM(result, i);
2200n/a PyTuple_SET_ITEM(result, i, elem);
2201n/a Py_DECREF(oldelem);
2202n/a break;
2203n/a }
2204n/a }
2205n/a
2206n/a /* If i is negative, then the indices have all rolled-over
2207n/a and we're done. */
2208n/a if (i < 0)
2209n/a goto empty;
2210n/a }
2211n/a
2212n/a Py_INCREF(result);
2213n/a return result;
2214n/a
2215n/aempty:
2216n/a lz->stopped = 1;
2217n/a return NULL;
2218n/a}
2219n/a
2220n/astatic PyObject *
2221n/aproduct_reduce(productobject *lz)
2222n/a{
2223n/a if (lz->stopped) {
2224n/a return Py_BuildValue("O(())", Py_TYPE(lz));
2225n/a } else if (lz->result == NULL) {
2226n/a return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2227n/a } else {
2228n/a PyObject *indices;
2229n/a Py_ssize_t n, i;
2230n/a
2231n/a /* we must pickle the indices use them for setstate, and
2232n/a * additionally indicate that the iterator has started
2233n/a */
2234n/a n = PyTuple_GET_SIZE(lz->pools);
2235n/a indices = PyTuple_New(n);
2236n/a if (indices == NULL)
2237n/a return NULL;
2238n/a for (i=0; i<n; i++){
2239n/a PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2240n/a if (!index) {
2241n/a Py_DECREF(indices);
2242n/a return NULL;
2243n/a }
2244n/a PyTuple_SET_ITEM(indices, i, index);
2245n/a }
2246n/a return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2247n/a }
2248n/a}
2249n/a
2250n/astatic PyObject *
2251n/aproduct_setstate(productobject *lz, PyObject *state)
2252n/a{
2253n/a PyObject *result;
2254n/a Py_ssize_t n, i;
2255n/a
2256n/a n = PyTuple_GET_SIZE(lz->pools);
2257n/a if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2258n/a PyErr_SetString(PyExc_ValueError, "invalid arguments");
2259n/a return NULL;
2260n/a }
2261n/a for (i=0; i<n; i++)
2262n/a {
2263n/a PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2264n/a Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2265n/a PyObject* pool;
2266n/a Py_ssize_t poolsize;
2267n/a if (index < 0 && PyErr_Occurred())
2268n/a return NULL; /* not an integer */
2269n/a pool = PyTuple_GET_ITEM(lz->pools, i);
2270n/a poolsize = PyTuple_GET_SIZE(pool);
2271n/a if (poolsize == 0) {
2272n/a lz->stopped = 1;
2273n/a Py_RETURN_NONE;
2274n/a }
2275n/a /* clamp the index */
2276n/a if (index < 0)
2277n/a index = 0;
2278n/a else if (index > poolsize-1)
2279n/a index = poolsize-1;
2280n/a lz->indices[i] = index;
2281n/a }
2282n/a
2283n/a result = PyTuple_New(n);
2284n/a if (!result)
2285n/a return NULL;
2286n/a for (i=0; i<n; i++) {
2287n/a PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2288n/a PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2289n/a Py_INCREF(element);
2290n/a PyTuple_SET_ITEM(result, i, element);
2291n/a }
2292n/a Py_XSETREF(lz->result, result);
2293n/a Py_RETURN_NONE;
2294n/a}
2295n/a
2296n/astatic PyMethodDef product_methods[] = {
2297n/a {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2298n/a reduce_doc},
2299n/a {"__setstate__", (PyCFunction)product_setstate, METH_O,
2300n/a setstate_doc},
2301n/a {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2302n/a sizeof_doc},
2303n/a {NULL, NULL} /* sentinel */
2304n/a};
2305n/a
2306n/aPyDoc_STRVAR(product_doc,
2307n/a"product(*iterables, repeat=1) --> product object\n\
2308n/a\n\
2309n/aCartesian product of input iterables. Equivalent to nested for-loops.\n\n\
2310n/aFor example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2311n/aThe leftmost iterators are in the outermost for-loop, so the output tuples\n\
2312n/acycle in a manner similar to an odometer (with the rightmost element changing\n\
2313n/aon every iteration).\n\n\
2314n/aTo compute the product of an iterable with itself, specify the number\n\
2315n/aof repetitions with the optional repeat keyword argument. For example,\n\
2316n/aproduct(A, repeat=4) means the same as product(A, A, A, A).\n\n\
2317n/aproduct('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2318n/aproduct((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2319n/a
2320n/astatic PyTypeObject product_type = {
2321n/a PyVarObject_HEAD_INIT(NULL, 0)
2322n/a "itertools.product", /* tp_name */
2323n/a sizeof(productobject), /* tp_basicsize */
2324n/a 0, /* tp_itemsize */
2325n/a /* methods */
2326n/a (destructor)product_dealloc, /* tp_dealloc */
2327n/a 0, /* tp_print */
2328n/a 0, /* tp_getattr */
2329n/a 0, /* tp_setattr */
2330n/a 0, /* tp_reserved */
2331n/a 0, /* tp_repr */
2332n/a 0, /* tp_as_number */
2333n/a 0, /* tp_as_sequence */
2334n/a 0, /* tp_as_mapping */
2335n/a 0, /* tp_hash */
2336n/a 0, /* tp_call */
2337n/a 0, /* tp_str */
2338n/a PyObject_GenericGetAttr, /* tp_getattro */
2339n/a 0, /* tp_setattro */
2340n/a 0, /* tp_as_buffer */
2341n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2342n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
2343n/a product_doc, /* tp_doc */
2344n/a (traverseproc)product_traverse, /* tp_traverse */
2345n/a 0, /* tp_clear */
2346n/a 0, /* tp_richcompare */
2347n/a 0, /* tp_weaklistoffset */
2348n/a PyObject_SelfIter, /* tp_iter */
2349n/a (iternextfunc)product_next, /* tp_iternext */
2350n/a product_methods, /* tp_methods */
2351n/a 0, /* tp_members */
2352n/a 0, /* tp_getset */
2353n/a 0, /* tp_base */
2354n/a 0, /* tp_dict */
2355n/a 0, /* tp_descr_get */
2356n/a 0, /* tp_descr_set */
2357n/a 0, /* tp_dictoffset */
2358n/a 0, /* tp_init */
2359n/a 0, /* tp_alloc */
2360n/a product_new, /* tp_new */
2361n/a PyObject_GC_Del, /* tp_free */
2362n/a};
2363n/a
2364n/a
2365n/a/* combinations object *******************************************************/
2366n/a
2367n/atypedef struct {
2368n/a PyObject_HEAD
2369n/a PyObject *pool; /* input converted to a tuple */
2370n/a Py_ssize_t *indices; /* one index per result element */
2371n/a PyObject *result; /* most recently returned result tuple */
2372n/a Py_ssize_t r; /* size of result tuple */
2373n/a int stopped; /* set to 1 when the iterator is exhausted */
2374n/a} combinationsobject;
2375n/a
2376n/astatic PyTypeObject combinations_type;
2377n/a
2378n/astatic PyObject *
2379n/acombinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2380n/a{
2381n/a combinationsobject *co;
2382n/a Py_ssize_t n;
2383n/a Py_ssize_t r;
2384n/a PyObject *pool = NULL;
2385n/a PyObject *iterable = NULL;
2386n/a Py_ssize_t *indices = NULL;
2387n/a Py_ssize_t i;
2388n/a static char *kwargs[] = {"iterable", "r", NULL};
2389n/a
2390n/a if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2391n/a &iterable, &r))
2392n/a return NULL;
2393n/a
2394n/a pool = PySequence_Tuple(iterable);
2395n/a if (pool == NULL)
2396n/a goto error;
2397n/a n = PyTuple_GET_SIZE(pool);
2398n/a if (r < 0) {
2399n/a PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2400n/a goto error;
2401n/a }
2402n/a
2403n/a indices = PyMem_New(Py_ssize_t, r);
2404n/a if (indices == NULL) {
2405n/a PyErr_NoMemory();
2406n/a goto error;
2407n/a }
2408n/a
2409n/a for (i=0 ; i<r ; i++)
2410n/a indices[i] = i;
2411n/a
2412n/a /* create combinationsobject structure */
2413n/a co = (combinationsobject *)type->tp_alloc(type, 0);
2414n/a if (co == NULL)
2415n/a goto error;
2416n/a
2417n/a co->pool = pool;
2418n/a co->indices = indices;
2419n/a co->result = NULL;
2420n/a co->r = r;
2421n/a co->stopped = r > n ? 1 : 0;
2422n/a
2423n/a return (PyObject *)co;
2424n/a
2425n/aerror:
2426n/a if (indices != NULL)
2427n/a PyMem_Free(indices);
2428n/a Py_XDECREF(pool);
2429n/a return NULL;
2430n/a}
2431n/a
2432n/astatic void
2433n/acombinations_dealloc(combinationsobject *co)
2434n/a{
2435n/a PyObject_GC_UnTrack(co);
2436n/a Py_XDECREF(co->pool);
2437n/a Py_XDECREF(co->result);
2438n/a if (co->indices != NULL)
2439n/a PyMem_Free(co->indices);
2440n/a Py_TYPE(co)->tp_free(co);
2441n/a}
2442n/a
2443n/astatic PyObject *
2444n/acombinations_sizeof(combinationsobject *co, void *unused)
2445n/a{
2446n/a Py_ssize_t res;
2447n/a
2448n/a res = _PyObject_SIZE(Py_TYPE(co));
2449n/a res += co->r * sizeof(Py_ssize_t);
2450n/a return PyLong_FromSsize_t(res);
2451n/a}
2452n/a
2453n/astatic int
2454n/acombinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2455n/a{
2456n/a Py_VISIT(co->pool);
2457n/a Py_VISIT(co->result);
2458n/a return 0;
2459n/a}
2460n/a
2461n/astatic PyObject *
2462n/acombinations_next(combinationsobject *co)
2463n/a{
2464n/a PyObject *elem;
2465n/a PyObject *oldelem;
2466n/a PyObject *pool = co->pool;
2467n/a Py_ssize_t *indices = co->indices;
2468n/a PyObject *result = co->result;
2469n/a Py_ssize_t n = PyTuple_GET_SIZE(pool);
2470n/a Py_ssize_t r = co->r;
2471n/a Py_ssize_t i, j, index;
2472n/a
2473n/a if (co->stopped)
2474n/a return NULL;
2475n/a
2476n/a if (result == NULL) {
2477n/a /* On the first pass, initialize result tuple using the indices */
2478n/a result = PyTuple_New(r);
2479n/a if (result == NULL)
2480n/a goto empty;
2481n/a co->result = result;
2482n/a for (i=0; i<r ; i++) {
2483n/a index = indices[i];
2484n/a elem = PyTuple_GET_ITEM(pool, index);
2485n/a Py_INCREF(elem);
2486n/a PyTuple_SET_ITEM(result, i, elem);
2487n/a }
2488n/a } else {
2489n/a /* Copy the previous result tuple or re-use it if available */
2490n/a if (Py_REFCNT(result) > 1) {
2491n/a PyObject *old_result = result;
2492n/a result = PyTuple_New(r);
2493n/a if (result == NULL)
2494n/a goto empty;
2495n/a co->result = result;
2496n/a for (i=0; i<r ; i++) {
2497n/a elem = PyTuple_GET_ITEM(old_result, i);
2498n/a Py_INCREF(elem);
2499n/a PyTuple_SET_ITEM(result, i, elem);
2500n/a }
2501n/a Py_DECREF(old_result);
2502n/a }
2503n/a /* Now, we've got the only copy so we can update it in-place
2504n/a * CPython's empty tuple is a singleton and cached in
2505n/a * PyTuple's freelist.
2506n/a */
2507n/a assert(r == 0 || Py_REFCNT(result) == 1);
2508n/a
2509n/a /* Scan indices right-to-left until finding one that is not
2510n/a at its maximum (i + n - r). */
2511n/a for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2512n/a ;
2513n/a
2514n/a /* If i is negative, then the indices are all at
2515n/a their maximum value and we're done. */
2516n/a if (i < 0)
2517n/a goto empty;
2518n/a
2519n/a /* Increment the current index which we know is not at its
2520n/a maximum. Then move back to the right setting each index
2521n/a to its lowest possible value (one higher than the index
2522n/a to its left -- this maintains the sort order invariant). */
2523n/a indices[i]++;
2524n/a for (j=i+1 ; j<r ; j++)
2525n/a indices[j] = indices[j-1] + 1;
2526n/a
2527n/a /* Update the result tuple for the new indices
2528n/a starting with i, the leftmost index that changed */
2529n/a for ( ; i<r ; i++) {
2530n/a index = indices[i];
2531n/a elem = PyTuple_GET_ITEM(pool, index);
2532n/a Py_INCREF(elem);
2533n/a oldelem = PyTuple_GET_ITEM(result, i);
2534n/a PyTuple_SET_ITEM(result, i, elem);
2535n/a Py_DECREF(oldelem);
2536n/a }
2537n/a }
2538n/a
2539n/a Py_INCREF(result);
2540n/a return result;
2541n/a
2542n/aempty:
2543n/a co->stopped = 1;
2544n/a return NULL;
2545n/a}
2546n/a
2547n/astatic PyObject *
2548n/acombinations_reduce(combinationsobject *lz)
2549n/a{
2550n/a if (lz->result == NULL) {
2551n/a return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2552n/a } else if (lz->stopped) {
2553n/a return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2554n/a } else {
2555n/a PyObject *indices;
2556n/a Py_ssize_t i;
2557n/a
2558n/a /* we must pickle the indices and use them for setstate */
2559n/a indices = PyTuple_New(lz->r);
2560n/a if (!indices)
2561n/a return NULL;
2562n/a for (i=0; i<lz->r; i++)
2563n/a {
2564n/a PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2565n/a if (!index) {
2566n/a Py_DECREF(indices);
2567n/a return NULL;
2568n/a }
2569n/a PyTuple_SET_ITEM(indices, i, index);
2570n/a }
2571n/a
2572n/a return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2573n/a }
2574n/a}
2575n/a
2576n/astatic PyObject *
2577n/acombinations_setstate(combinationsobject *lz, PyObject *state)
2578n/a{
2579n/a PyObject *result;
2580n/a Py_ssize_t i;
2581n/a Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2582n/a
2583n/a if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
2584n/a PyErr_SetString(PyExc_ValueError, "invalid arguments");
2585n/a return NULL;
2586n/a }
2587n/a
2588n/a for (i=0; i<lz->r; i++) {
2589n/a Py_ssize_t max;
2590n/a PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2591n/a Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2592n/a
2593n/a if (index == -1 && PyErr_Occurred())
2594n/a return NULL; /* not an integer */
2595n/a max = i + n - lz->r;
2596n/a /* clamp the index (beware of negative max) */
2597n/a if (index > max)
2598n/a index = max;
2599n/a if (index < 0)
2600n/a index = 0;
2601n/a lz->indices[i] = index;
2602n/a }
2603n/a
2604n/a result = PyTuple_New(lz->r);
2605n/a if (result == NULL)
2606n/a return NULL;
2607n/a for (i=0; i<lz->r; i++) {
2608n/a PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2609n/a Py_INCREF(element);
2610n/a PyTuple_SET_ITEM(result, i, element);
2611n/a }
2612n/a
2613n/a Py_XSETREF(lz->result, result);
2614n/a Py_RETURN_NONE;
2615n/a}
2616n/a
2617n/astatic PyMethodDef combinations_methods[] = {
2618n/a {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2619n/a reduce_doc},
2620n/a {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2621n/a setstate_doc},
2622n/a {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2623n/a sizeof_doc},
2624n/a {NULL, NULL} /* sentinel */
2625n/a};
2626n/a
2627n/aPyDoc_STRVAR(combinations_doc,
2628n/a"combinations(iterable, r) --> combinations object\n\
2629n/a\n\
2630n/aReturn successive r-length combinations of elements in the iterable.\n\n\
2631n/acombinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2632n/a
2633n/astatic PyTypeObject combinations_type = {
2634n/a PyVarObject_HEAD_INIT(NULL, 0)
2635n/a "itertools.combinations", /* tp_name */
2636n/a sizeof(combinationsobject), /* tp_basicsize */
2637n/a 0, /* tp_itemsize */
2638n/a /* methods */
2639n/a (destructor)combinations_dealloc, /* tp_dealloc */
2640n/a 0, /* tp_print */
2641n/a 0, /* tp_getattr */
2642n/a 0, /* tp_setattr */
2643n/a 0, /* tp_reserved */
2644n/a 0, /* tp_repr */
2645n/a 0, /* tp_as_number */
2646n/a 0, /* tp_as_sequence */
2647n/a 0, /* tp_as_mapping */
2648n/a 0, /* tp_hash */
2649n/a 0, /* tp_call */
2650n/a 0, /* tp_str */
2651n/a PyObject_GenericGetAttr, /* tp_getattro */
2652n/a 0, /* tp_setattro */
2653n/a 0, /* tp_as_buffer */
2654n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2655n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
2656n/a combinations_doc, /* tp_doc */
2657n/a (traverseproc)combinations_traverse,/* tp_traverse */
2658n/a 0, /* tp_clear */
2659n/a 0, /* tp_richcompare */
2660n/a 0, /* tp_weaklistoffset */
2661n/a PyObject_SelfIter, /* tp_iter */
2662n/a (iternextfunc)combinations_next, /* tp_iternext */
2663n/a combinations_methods, /* tp_methods */
2664n/a 0, /* tp_members */
2665n/a 0, /* tp_getset */
2666n/a 0, /* tp_base */
2667n/a 0, /* tp_dict */
2668n/a 0, /* tp_descr_get */
2669n/a 0, /* tp_descr_set */
2670n/a 0, /* tp_dictoffset */
2671n/a 0, /* tp_init */
2672n/a 0, /* tp_alloc */
2673n/a combinations_new, /* tp_new */
2674n/a PyObject_GC_Del, /* tp_free */
2675n/a};
2676n/a
2677n/a
2678n/a/* combinations with replacement object **************************************/
2679n/a
2680n/a/* Equivalent to:
2681n/a
2682n/a def combinations_with_replacement(iterable, r):
2683n/a "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2684n/a # number items returned: (n+r-1)! / r! / (n-1)!
2685n/a pool = tuple(iterable)
2686n/a n = len(pool)
2687n/a indices = [0] * r
2688n/a yield tuple(pool[i] for i in indices)
2689n/a while 1:
2690n/a for i in reversed(range(r)):
2691n/a if indices[i] != n - 1:
2692n/a break
2693n/a else:
2694n/a return
2695n/a indices[i:] = [indices[i] + 1] * (r - i)
2696n/a yield tuple(pool[i] for i in indices)
2697n/a
2698n/a def combinations_with_replacement2(iterable, r):
2699n/a 'Alternate version that filters from product()'
2700n/a pool = tuple(iterable)
2701n/a n = len(pool)
2702n/a for indices in product(range(n), repeat=r):
2703n/a if sorted(indices) == list(indices):
2704n/a yield tuple(pool[i] for i in indices)
2705n/a*/
2706n/atypedef struct {
2707n/a PyObject_HEAD
2708n/a PyObject *pool; /* input converted to a tuple */
2709n/a Py_ssize_t *indices; /* one index per result element */
2710n/a PyObject *result; /* most recently returned result tuple */
2711n/a Py_ssize_t r; /* size of result tuple */
2712n/a int stopped; /* set to 1 when the cwr iterator is exhausted */
2713n/a} cwrobject;
2714n/a
2715n/astatic PyTypeObject cwr_type;
2716n/a
2717n/astatic PyObject *
2718n/acwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2719n/a{
2720n/a cwrobject *co;
2721n/a Py_ssize_t n;
2722n/a Py_ssize_t r;
2723n/a PyObject *pool = NULL;
2724n/a PyObject *iterable = NULL;
2725n/a Py_ssize_t *indices = NULL;
2726n/a Py_ssize_t i;
2727n/a static char *kwargs[] = {"iterable", "r", NULL};
2728n/a
2729n/a if (!PyArg_ParseTupleAndKeywords(args, kwds,
2730n/a "On:combinations_with_replacement",
2731n/a kwargs, &iterable, &r))
2732n/a return NULL;
2733n/a
2734n/a pool = PySequence_Tuple(iterable);
2735n/a if (pool == NULL)
2736n/a goto error;
2737n/a n = PyTuple_GET_SIZE(pool);
2738n/a if (r < 0) {
2739n/a PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2740n/a goto error;
2741n/a }
2742n/a
2743n/a indices = PyMem_New(Py_ssize_t, r);
2744n/a if (indices == NULL) {
2745n/a PyErr_NoMemory();
2746n/a goto error;
2747n/a }
2748n/a
2749n/a for (i=0 ; i<r ; i++)
2750n/a indices[i] = 0;
2751n/a
2752n/a /* create cwrobject structure */
2753n/a co = (cwrobject *)type->tp_alloc(type, 0);
2754n/a if (co == NULL)
2755n/a goto error;
2756n/a
2757n/a co->pool = pool;
2758n/a co->indices = indices;
2759n/a co->result = NULL;
2760n/a co->r = r;
2761n/a co->stopped = !n && r;
2762n/a
2763n/a return (PyObject *)co;
2764n/a
2765n/aerror:
2766n/a if (indices != NULL)
2767n/a PyMem_Free(indices);
2768n/a Py_XDECREF(pool);
2769n/a return NULL;
2770n/a}
2771n/a
2772n/astatic void
2773n/acwr_dealloc(cwrobject *co)
2774n/a{
2775n/a PyObject_GC_UnTrack(co);
2776n/a Py_XDECREF(co->pool);
2777n/a Py_XDECREF(co->result);
2778n/a if (co->indices != NULL)
2779n/a PyMem_Free(co->indices);
2780n/a Py_TYPE(co)->tp_free(co);
2781n/a}
2782n/a
2783n/astatic PyObject *
2784n/acwr_sizeof(cwrobject *co, void *unused)
2785n/a{
2786n/a Py_ssize_t res;
2787n/a
2788n/a res = _PyObject_SIZE(Py_TYPE(co));
2789n/a res += co->r * sizeof(Py_ssize_t);
2790n/a return PyLong_FromSsize_t(res);
2791n/a}
2792n/a
2793n/astatic int
2794n/acwr_traverse(cwrobject *co, visitproc visit, void *arg)
2795n/a{
2796n/a Py_VISIT(co->pool);
2797n/a Py_VISIT(co->result);
2798n/a return 0;
2799n/a}
2800n/a
2801n/astatic PyObject *
2802n/acwr_next(cwrobject *co)
2803n/a{
2804n/a PyObject *elem;
2805n/a PyObject *oldelem;
2806n/a PyObject *pool = co->pool;
2807n/a Py_ssize_t *indices = co->indices;
2808n/a PyObject *result = co->result;
2809n/a Py_ssize_t n = PyTuple_GET_SIZE(pool);
2810n/a Py_ssize_t r = co->r;
2811n/a Py_ssize_t i, index;
2812n/a
2813n/a if (co->stopped)
2814n/a return NULL;
2815n/a
2816n/a if (result == NULL) {
2817n/a /* On the first pass, initialize result tuple with pool[0] */
2818n/a result = PyTuple_New(r);
2819n/a if (result == NULL)
2820n/a goto empty;
2821n/a co->result = result;
2822n/a if (n > 0) {
2823n/a elem = PyTuple_GET_ITEM(pool, 0);
2824n/a for (i=0; i<r ; i++) {
2825n/a assert(indices[i] == 0);
2826n/a Py_INCREF(elem);
2827n/a PyTuple_SET_ITEM(result, i, elem);
2828n/a }
2829n/a }
2830n/a } else {
2831n/a /* Copy the previous result tuple or re-use it if available */
2832n/a if (Py_REFCNT(result) > 1) {
2833n/a PyObject *old_result = result;
2834n/a result = PyTuple_New(r);
2835n/a if (result == NULL)
2836n/a goto empty;
2837n/a co->result = result;
2838n/a for (i=0; i<r ; i++) {
2839n/a elem = PyTuple_GET_ITEM(old_result, i);
2840n/a Py_INCREF(elem);
2841n/a PyTuple_SET_ITEM(result, i, elem);
2842n/a }
2843n/a Py_DECREF(old_result);
2844n/a }
2845n/a /* Now, we've got the only copy so we can update it in-place CPython's
2846n/a empty tuple is a singleton and cached in PyTuple's freelist. */
2847n/a assert(r == 0 || Py_REFCNT(result) == 1);
2848n/a
2849n/a /* Scan indices right-to-left until finding one that is not
2850n/a * at its maximum (n-1). */
2851n/a for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2852n/a ;
2853n/a
2854n/a /* If i is negative, then the indices are all at
2855n/a their maximum value and we're done. */
2856n/a if (i < 0)
2857n/a goto empty;
2858n/a
2859n/a /* Increment the current index which we know is not at its
2860n/a maximum. Then set all to the right to the same value. */
2861n/a index = indices[i] + 1;
2862n/a assert(index < n);
2863n/a elem = PyTuple_GET_ITEM(pool, index);
2864n/a for ( ; i<r ; i++) {
2865n/a indices[i] = index;
2866n/a Py_INCREF(elem);
2867n/a oldelem = PyTuple_GET_ITEM(result, i);
2868n/a PyTuple_SET_ITEM(result, i, elem);
2869n/a Py_DECREF(oldelem);
2870n/a }
2871n/a }
2872n/a
2873n/a Py_INCREF(result);
2874n/a return result;
2875n/a
2876n/aempty:
2877n/a co->stopped = 1;
2878n/a return NULL;
2879n/a}
2880n/a
2881n/astatic PyObject *
2882n/acwr_reduce(cwrobject *lz)
2883n/a{
2884n/a if (lz->result == NULL) {
2885n/a return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2886n/a } else if (lz->stopped) {
2887n/a return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2888n/a } else {
2889n/a PyObject *indices;
2890n/a Py_ssize_t i;
2891n/a
2892n/a /* we must pickle the indices and use them for setstate */
2893n/a indices = PyTuple_New(lz->r);
2894n/a if (!indices)
2895n/a return NULL;
2896n/a for (i=0; i<lz->r; i++) {
2897n/a PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2898n/a if (!index) {
2899n/a Py_DECREF(indices);
2900n/a return NULL;
2901n/a }
2902n/a PyTuple_SET_ITEM(indices, i, index);
2903n/a }
2904n/a
2905n/a return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2906n/a }
2907n/a}
2908n/a
2909n/astatic PyObject *
2910n/acwr_setstate(cwrobject *lz, PyObject *state)
2911n/a{
2912n/a PyObject *result;
2913n/a Py_ssize_t n, i;
2914n/a
2915n/a if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2916n/a {
2917n/a PyErr_SetString(PyExc_ValueError, "invalid arguments");
2918n/a return NULL;
2919n/a }
2920n/a
2921n/a n = PyTuple_GET_SIZE(lz->pool);
2922n/a for (i=0; i<lz->r; i++) {
2923n/a PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2924n/a Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2925n/a
2926n/a if (index < 0 && PyErr_Occurred())
2927n/a return NULL; /* not an integer */
2928n/a /* clamp the index */
2929n/a if (index < 0)
2930n/a index = 0;
2931n/a else if (index > n-1)
2932n/a index = n-1;
2933n/a lz->indices[i] = index;
2934n/a }
2935n/a result = PyTuple_New(lz->r);
2936n/a if (result == NULL)
2937n/a return NULL;
2938n/a for (i=0; i<lz->r; i++) {
2939n/a PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2940n/a Py_INCREF(element);
2941n/a PyTuple_SET_ITEM(result, i, element);
2942n/a }
2943n/a Py_XSETREF(lz->result, result);
2944n/a Py_RETURN_NONE;
2945n/a}
2946n/a
2947n/astatic PyMethodDef cwr_methods[] = {
2948n/a {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2949n/a reduce_doc},
2950n/a {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2951n/a setstate_doc},
2952n/a {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2953n/a sizeof_doc},
2954n/a {NULL, NULL} /* sentinel */
2955n/a};
2956n/a
2957n/aPyDoc_STRVAR(cwr_doc,
2958n/a"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
2959n/a\n\
2960n/aReturn successive r-length combinations of elements in the iterable\n\
2961n/aallowing individual elements to have successive repeats.\n\
2962n/acombinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2963n/a
2964n/astatic PyTypeObject cwr_type = {
2965n/a PyVarObject_HEAD_INIT(NULL, 0)
2966n/a "itertools.combinations_with_replacement", /* tp_name */
2967n/a sizeof(cwrobject), /* tp_basicsize */
2968n/a 0, /* tp_itemsize */
2969n/a /* methods */
2970n/a (destructor)cwr_dealloc, /* tp_dealloc */
2971n/a 0, /* tp_print */
2972n/a 0, /* tp_getattr */
2973n/a 0, /* tp_setattr */
2974n/a 0, /* tp_reserved */
2975n/a 0, /* tp_repr */
2976n/a 0, /* tp_as_number */
2977n/a 0, /* tp_as_sequence */
2978n/a 0, /* tp_as_mapping */
2979n/a 0, /* tp_hash */
2980n/a 0, /* tp_call */
2981n/a 0, /* tp_str */
2982n/a PyObject_GenericGetAttr, /* tp_getattro */
2983n/a 0, /* tp_setattro */
2984n/a 0, /* tp_as_buffer */
2985n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2986n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
2987n/a cwr_doc, /* tp_doc */
2988n/a (traverseproc)cwr_traverse, /* tp_traverse */
2989n/a 0, /* tp_clear */
2990n/a 0, /* tp_richcompare */
2991n/a 0, /* tp_weaklistoffset */
2992n/a PyObject_SelfIter, /* tp_iter */
2993n/a (iternextfunc)cwr_next, /* tp_iternext */
2994n/a cwr_methods, /* tp_methods */
2995n/a 0, /* tp_members */
2996n/a 0, /* tp_getset */
2997n/a 0, /* tp_base */
2998n/a 0, /* tp_dict */
2999n/a 0, /* tp_descr_get */
3000n/a 0, /* tp_descr_set */
3001n/a 0, /* tp_dictoffset */
3002n/a 0, /* tp_init */
3003n/a 0, /* tp_alloc */
3004n/a cwr_new, /* tp_new */
3005n/a PyObject_GC_Del, /* tp_free */
3006n/a};
3007n/a
3008n/a
3009n/a/* permutations object ********************************************************
3010n/a
3011n/adef permutations(iterable, r=None):
3012n/a 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3013n/a pool = tuple(iterable)
3014n/a n = len(pool)
3015n/a r = n if r is None else r
3016n/a indices = range(n)
3017n/a cycles = range(n-r+1, n+1)[::-1]
3018n/a yield tuple(pool[i] for i in indices[:r])
3019n/a while n:
3020n/a for i in reversed(range(r)):
3021n/a cycles[i] -= 1
3022n/a if cycles[i] == 0:
3023n/a indices[i:] = indices[i+1:] + indices[i:i+1]
3024n/a cycles[i] = n - i
3025n/a else:
3026n/a j = cycles[i]
3027n/a indices[i], indices[-j] = indices[-j], indices[i]
3028n/a yield tuple(pool[i] for i in indices[:r])
3029n/a break
3030n/a else:
3031n/a return
3032n/a*/
3033n/a
3034n/atypedef struct {
3035n/a PyObject_HEAD
3036n/a PyObject *pool; /* input converted to a tuple */
3037n/a Py_ssize_t *indices; /* one index per element in the pool */
3038n/a Py_ssize_t *cycles; /* one rollover counter per element in the result */
3039n/a PyObject *result; /* most recently returned result tuple */
3040n/a Py_ssize_t r; /* size of result tuple */
3041n/a int stopped; /* set to 1 when the iterator is exhausted */
3042n/a} permutationsobject;
3043n/a
3044n/astatic PyTypeObject permutations_type;
3045n/a
3046n/astatic PyObject *
3047n/apermutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3048n/a{
3049n/a permutationsobject *po;
3050n/a Py_ssize_t n;
3051n/a Py_ssize_t r;
3052n/a PyObject *robj = Py_None;
3053n/a PyObject *pool = NULL;
3054n/a PyObject *iterable = NULL;
3055n/a Py_ssize_t *indices = NULL;
3056n/a Py_ssize_t *cycles = NULL;
3057n/a Py_ssize_t i;
3058n/a static char *kwargs[] = {"iterable", "r", NULL};
3059n/a
3060n/a if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3061n/a &iterable, &robj))
3062n/a return NULL;
3063n/a
3064n/a pool = PySequence_Tuple(iterable);
3065n/a if (pool == NULL)
3066n/a goto error;
3067n/a n = PyTuple_GET_SIZE(pool);
3068n/a
3069n/a r = n;
3070n/a if (robj != Py_None) {
3071n/a if (!PyLong_Check(robj)) {
3072n/a PyErr_SetString(PyExc_TypeError, "Expected int as r");
3073n/a goto error;
3074n/a }
3075n/a r = PyLong_AsSsize_t(robj);
3076n/a if (r == -1 && PyErr_Occurred())
3077n/a goto error;
3078n/a }
3079n/a if (r < 0) {
3080n/a PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3081n/a goto error;
3082n/a }
3083n/a
3084n/a indices = PyMem_New(Py_ssize_t, n);
3085n/a cycles = PyMem_New(Py_ssize_t, r);
3086n/a if (indices == NULL || cycles == NULL) {
3087n/a PyErr_NoMemory();
3088n/a goto error;
3089n/a }
3090n/a
3091n/a for (i=0 ; i<n ; i++)
3092n/a indices[i] = i;
3093n/a for (i=0 ; i<r ; i++)
3094n/a cycles[i] = n - i;
3095n/a
3096n/a /* create permutationsobject structure */
3097n/a po = (permutationsobject *)type->tp_alloc(type, 0);
3098n/a if (po == NULL)
3099n/a goto error;
3100n/a
3101n/a po->pool = pool;
3102n/a po->indices = indices;
3103n/a po->cycles = cycles;
3104n/a po->result = NULL;
3105n/a po->r = r;
3106n/a po->stopped = r > n ? 1 : 0;
3107n/a
3108n/a return (PyObject *)po;
3109n/a
3110n/aerror:
3111n/a if (indices != NULL)
3112n/a PyMem_Free(indices);
3113n/a if (cycles != NULL)
3114n/a PyMem_Free(cycles);
3115n/a Py_XDECREF(pool);
3116n/a return NULL;
3117n/a}
3118n/a
3119n/astatic void
3120n/apermutations_dealloc(permutationsobject *po)
3121n/a{
3122n/a PyObject_GC_UnTrack(po);
3123n/a Py_XDECREF(po->pool);
3124n/a Py_XDECREF(po->result);
3125n/a PyMem_Free(po->indices);
3126n/a PyMem_Free(po->cycles);
3127n/a Py_TYPE(po)->tp_free(po);
3128n/a}
3129n/a
3130n/astatic PyObject *
3131n/apermutations_sizeof(permutationsobject *po, void *unused)
3132n/a{
3133n/a Py_ssize_t res;
3134n/a
3135n/a res = _PyObject_SIZE(Py_TYPE(po));
3136n/a res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3137n/a res += po->r * sizeof(Py_ssize_t);
3138n/a return PyLong_FromSsize_t(res);
3139n/a}
3140n/a
3141n/astatic int
3142n/apermutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3143n/a{
3144n/a Py_VISIT(po->pool);
3145n/a Py_VISIT(po->result);
3146n/a return 0;
3147n/a}
3148n/a
3149n/astatic PyObject *
3150n/apermutations_next(permutationsobject *po)
3151n/a{
3152n/a PyObject *elem;
3153n/a PyObject *oldelem;
3154n/a PyObject *pool = po->pool;
3155n/a Py_ssize_t *indices = po->indices;
3156n/a Py_ssize_t *cycles = po->cycles;
3157n/a PyObject *result = po->result;
3158n/a Py_ssize_t n = PyTuple_GET_SIZE(pool);
3159n/a Py_ssize_t r = po->r;
3160n/a Py_ssize_t i, j, k, index;
3161n/a
3162n/a if (po->stopped)
3163n/a return NULL;
3164n/a
3165n/a if (result == NULL) {
3166n/a /* On the first pass, initialize result tuple using the indices */
3167n/a result = PyTuple_New(r);
3168n/a if (result == NULL)
3169n/a goto empty;
3170n/a po->result = result;
3171n/a for (i=0; i<r ; i++) {
3172n/a index = indices[i];
3173n/a elem = PyTuple_GET_ITEM(pool, index);
3174n/a Py_INCREF(elem);
3175n/a PyTuple_SET_ITEM(result, i, elem);
3176n/a }
3177n/a } else {
3178n/a if (n == 0)
3179n/a goto empty;
3180n/a
3181n/a /* Copy the previous result tuple or re-use it if available */
3182n/a if (Py_REFCNT(result) > 1) {
3183n/a PyObject *old_result = result;
3184n/a result = PyTuple_New(r);
3185n/a if (result == NULL)
3186n/a goto empty;
3187n/a po->result = result;
3188n/a for (i=0; i<r ; i++) {
3189n/a elem = PyTuple_GET_ITEM(old_result, i);
3190n/a Py_INCREF(elem);
3191n/a PyTuple_SET_ITEM(result, i, elem);
3192n/a }
3193n/a Py_DECREF(old_result);
3194n/a }
3195n/a /* Now, we've got the only copy so we can update it in-place */
3196n/a assert(r == 0 || Py_REFCNT(result) == 1);
3197n/a
3198n/a /* Decrement rightmost cycle, moving leftward upon zero rollover */
3199n/a for (i=r-1 ; i>=0 ; i--) {
3200n/a cycles[i] -= 1;
3201n/a if (cycles[i] == 0) {
3202n/a /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3203n/a index = indices[i];
3204n/a for (j=i ; j<n-1 ; j++)
3205n/a indices[j] = indices[j+1];
3206n/a indices[n-1] = index;
3207n/a cycles[i] = n - i;
3208n/a } else {
3209n/a j = cycles[i];
3210n/a index = indices[i];
3211n/a indices[i] = indices[n-j];
3212n/a indices[n-j] = index;
3213n/a
3214n/a for (k=i; k<r ; k++) {
3215n/a /* start with i, the leftmost element that changed */
3216n/a /* yield tuple(pool[k] for k in indices[:r]) */
3217n/a index = indices[k];
3218n/a elem = PyTuple_GET_ITEM(pool, index);
3219n/a Py_INCREF(elem);
3220n/a oldelem = PyTuple_GET_ITEM(result, k);
3221n/a PyTuple_SET_ITEM(result, k, elem);
3222n/a Py_DECREF(oldelem);
3223n/a }
3224n/a break;
3225n/a }
3226n/a }
3227n/a /* If i is negative, then the cycles have all
3228n/a rolled-over and we're done. */
3229n/a if (i < 0)
3230n/a goto empty;
3231n/a }
3232n/a Py_INCREF(result);
3233n/a return result;
3234n/a
3235n/aempty:
3236n/a po->stopped = 1;
3237n/a return NULL;
3238n/a}
3239n/a
3240n/astatic PyObject *
3241n/apermutations_reduce(permutationsobject *po)
3242n/a{
3243n/a if (po->result == NULL) {
3244n/a return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3245n/a } else if (po->stopped) {
3246n/a return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3247n/a } else {
3248n/a PyObject *indices=NULL, *cycles=NULL;
3249n/a Py_ssize_t n, i;
3250n/a
3251n/a /* we must pickle the indices and cycles and use them for setstate */
3252n/a n = PyTuple_GET_SIZE(po->pool);
3253n/a indices = PyTuple_New(n);
3254n/a if (indices == NULL)
3255n/a goto err;
3256n/a for (i=0; i<n; i++) {
3257n/a PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3258n/a if (!index)
3259n/a goto err;
3260n/a PyTuple_SET_ITEM(indices, i, index);
3261n/a }
3262n/a
3263n/a cycles = PyTuple_New(po->r);
3264n/a if (cycles == NULL)
3265n/a goto err;
3266n/a for (i=0 ; i<po->r ; i++) {
3267n/a PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3268n/a if (!index)
3269n/a goto err;
3270n/a PyTuple_SET_ITEM(cycles, i, index);
3271n/a }
3272n/a return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3273n/a po->pool, po->r,
3274n/a indices, cycles);
3275n/a err:
3276n/a Py_XDECREF(indices);
3277n/a Py_XDECREF(cycles);
3278n/a return NULL;
3279n/a }
3280n/a}
3281n/a
3282n/astatic PyObject *
3283n/apermutations_setstate(permutationsobject *po, PyObject *state)
3284n/a{
3285n/a PyObject *indices, *cycles, *result;
3286n/a Py_ssize_t n, i;
3287n/a
3288n/a if (!PyTuple_Check(state)) {
3289n/a PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3290n/a return NULL;
3291n/a }
3292n/a if (!PyArg_ParseTuple(state, "O!O!",
3293n/a &PyTuple_Type, &indices,
3294n/a &PyTuple_Type, &cycles)) {
3295n/a return NULL;
3296n/a }
3297n/a
3298n/a n = PyTuple_GET_SIZE(po->pool);
3299n/a if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
3300n/a PyErr_SetString(PyExc_ValueError, "invalid arguments");
3301n/a return NULL;
3302n/a }
3303n/a
3304n/a for (i=0; i<n; i++) {
3305n/a PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3306n/a Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3307n/a if (index < 0 && PyErr_Occurred())
3308n/a return NULL; /* not an integer */
3309n/a /* clamp the index */
3310n/a if (index < 0)
3311n/a index = 0;
3312n/a else if (index > n-1)
3313n/a index = n-1;
3314n/a po->indices[i] = index;
3315n/a }
3316n/a
3317n/a for (i=0; i<po->r; i++) {
3318n/a PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3319n/a Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3320n/a if (index < 0 && PyErr_Occurred())
3321n/a return NULL; /* not an integer */
3322n/a if (index < 1)
3323n/a index = 1;
3324n/a else if (index > n-i)
3325n/a index = n-i;
3326n/a po->cycles[i] = index;
3327n/a }
3328n/a result = PyTuple_New(po->r);
3329n/a if (result == NULL)
3330n/a return NULL;
3331n/a for (i=0; i<po->r; i++) {
3332n/a PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3333n/a Py_INCREF(element);
3334n/a PyTuple_SET_ITEM(result, i, element);
3335n/a }
3336n/a Py_XSETREF(po->result, result);
3337n/a Py_RETURN_NONE;
3338n/a}
3339n/a
3340n/astatic PyMethodDef permuations_methods[] = {
3341n/a {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3342n/a reduce_doc},
3343n/a {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3344n/a setstate_doc},
3345n/a {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3346n/a sizeof_doc},
3347n/a {NULL, NULL} /* sentinel */
3348n/a};
3349n/a
3350n/aPyDoc_STRVAR(permutations_doc,
3351n/a"permutations(iterable[, r]) --> permutations object\n\
3352n/a\n\
3353n/aReturn successive r-length permutations of elements in the iterable.\n\n\
3354n/apermutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
3355n/a
3356n/astatic PyTypeObject permutations_type = {
3357n/a PyVarObject_HEAD_INIT(NULL, 0)
3358n/a "itertools.permutations", /* tp_name */
3359n/a sizeof(permutationsobject), /* tp_basicsize */
3360n/a 0, /* tp_itemsize */
3361n/a /* methods */
3362n/a (destructor)permutations_dealloc, /* tp_dealloc */
3363n/a 0, /* tp_print */
3364n/a 0, /* tp_getattr */
3365n/a 0, /* tp_setattr */
3366n/a 0, /* tp_reserved */
3367n/a 0, /* tp_repr */
3368n/a 0, /* tp_as_number */
3369n/a 0, /* tp_as_sequence */
3370n/a 0, /* tp_as_mapping */
3371n/a 0, /* tp_hash */
3372n/a 0, /* tp_call */
3373n/a 0, /* tp_str */
3374n/a PyObject_GenericGetAttr, /* tp_getattro */
3375n/a 0, /* tp_setattro */
3376n/a 0, /* tp_as_buffer */
3377n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3378n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
3379n/a permutations_doc, /* tp_doc */
3380n/a (traverseproc)permutations_traverse,/* tp_traverse */
3381n/a 0, /* tp_clear */
3382n/a 0, /* tp_richcompare */
3383n/a 0, /* tp_weaklistoffset */
3384n/a PyObject_SelfIter, /* tp_iter */
3385n/a (iternextfunc)permutations_next, /* tp_iternext */
3386n/a permuations_methods, /* tp_methods */
3387n/a 0, /* tp_members */
3388n/a 0, /* tp_getset */
3389n/a 0, /* tp_base */
3390n/a 0, /* tp_dict */
3391n/a 0, /* tp_descr_get */
3392n/a 0, /* tp_descr_set */
3393n/a 0, /* tp_dictoffset */
3394n/a 0, /* tp_init */
3395n/a 0, /* tp_alloc */
3396n/a permutations_new, /* tp_new */
3397n/a PyObject_GC_Del, /* tp_free */
3398n/a};
3399n/a
3400n/a/* accumulate object ********************************************************/
3401n/a
3402n/atypedef struct {
3403n/a PyObject_HEAD
3404n/a PyObject *total;
3405n/a PyObject *it;
3406n/a PyObject *binop;
3407n/a} accumulateobject;
3408n/a
3409n/astatic PyTypeObject accumulate_type;
3410n/a
3411n/astatic PyObject *
3412n/aaccumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3413n/a{
3414n/a static char *kwargs[] = {"iterable", "func", NULL};
3415n/a PyObject *iterable;
3416n/a PyObject *it;
3417n/a PyObject *binop = Py_None;
3418n/a accumulateobject *lz;
3419n/a
3420n/a if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3421n/a kwargs, &iterable, &binop))
3422n/a return NULL;
3423n/a
3424n/a /* Get iterator. */
3425n/a it = PyObject_GetIter(iterable);
3426n/a if (it == NULL)
3427n/a return NULL;
3428n/a
3429n/a /* create accumulateobject structure */
3430n/a lz = (accumulateobject *)type->tp_alloc(type, 0);
3431n/a if (lz == NULL) {
3432n/a Py_DECREF(it);
3433n/a return NULL;
3434n/a }
3435n/a
3436n/a if (binop != Py_None) {
3437n/a Py_XINCREF(binop);
3438n/a lz->binop = binop;
3439n/a }
3440n/a lz->total = NULL;
3441n/a lz->it = it;
3442n/a return (PyObject *)lz;
3443n/a}
3444n/a
3445n/astatic void
3446n/aaccumulate_dealloc(accumulateobject *lz)
3447n/a{
3448n/a PyObject_GC_UnTrack(lz);
3449n/a Py_XDECREF(lz->binop);
3450n/a Py_XDECREF(lz->total);
3451n/a Py_XDECREF(lz->it);
3452n/a Py_TYPE(lz)->tp_free(lz);
3453n/a}
3454n/a
3455n/astatic int
3456n/aaccumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3457n/a{
3458n/a Py_VISIT(lz->binop);
3459n/a Py_VISIT(lz->it);
3460n/a Py_VISIT(lz->total);
3461n/a return 0;
3462n/a}
3463n/a
3464n/astatic PyObject *
3465n/aaccumulate_next(accumulateobject *lz)
3466n/a{
3467n/a PyObject *val, *newtotal;
3468n/a
3469n/a val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
3470n/a if (val == NULL)
3471n/a return NULL;
3472n/a
3473n/a if (lz->total == NULL) {
3474n/a Py_INCREF(val);
3475n/a lz->total = val;
3476n/a return lz->total;
3477n/a }
3478n/a
3479n/a if (lz->binop == NULL)
3480n/a newtotal = PyNumber_Add(lz->total, val);
3481n/a else
3482n/a newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
3483n/a Py_DECREF(val);
3484n/a if (newtotal == NULL)
3485n/a return NULL;
3486n/a
3487n/a Py_INCREF(newtotal);
3488n/a Py_SETREF(lz->total, newtotal);
3489n/a return newtotal;
3490n/a}
3491n/a
3492n/astatic PyObject *
3493n/aaccumulate_reduce(accumulateobject *lz)
3494n/a{
3495n/a if (lz->total == Py_None) {
3496n/a PyObject *it;
3497n/a
3498n/a if (PyType_Ready(&chain_type) < 0)
3499n/a return NULL;
3500n/a if (PyType_Ready(&islice_type) < 0)
3501n/a return NULL;
3502n/a it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3503n/a lz->total, lz->it);
3504n/a if (it == NULL)
3505n/a return NULL;
3506n/a it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3507n/a it, lz->binop ? lz->binop : Py_None);
3508n/a if (it == NULL)
3509n/a return NULL;
3510n/a return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3511n/a }
3512n/a return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3513n/a lz->it, lz->binop?lz->binop:Py_None,
3514n/a lz->total?lz->total:Py_None);
3515n/a}
3516n/a
3517n/astatic PyObject *
3518n/aaccumulate_setstate(accumulateobject *lz, PyObject *state)
3519n/a{
3520n/a Py_INCREF(state);
3521n/a Py_XSETREF(lz->total, state);
3522n/a Py_RETURN_NONE;
3523n/a}
3524n/a
3525n/astatic PyMethodDef accumulate_methods[] = {
3526n/a {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3527n/a reduce_doc},
3528n/a {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3529n/a setstate_doc},
3530n/a {NULL, NULL} /* sentinel */
3531n/a};
3532n/a
3533n/aPyDoc_STRVAR(accumulate_doc,
3534n/a"accumulate(iterable[, func]) --> accumulate object\n\
3535n/a\n\
3536n/aReturn series of accumulated sums (or other binary function results).");
3537n/a
3538n/astatic PyTypeObject accumulate_type = {
3539n/a PyVarObject_HEAD_INIT(NULL, 0)
3540n/a "itertools.accumulate", /* tp_name */
3541n/a sizeof(accumulateobject), /* tp_basicsize */
3542n/a 0, /* tp_itemsize */
3543n/a /* methods */
3544n/a (destructor)accumulate_dealloc, /* tp_dealloc */
3545n/a 0, /* tp_print */
3546n/a 0, /* tp_getattr */
3547n/a 0, /* tp_setattr */
3548n/a 0, /* tp_reserved */
3549n/a 0, /* tp_repr */
3550n/a 0, /* tp_as_number */
3551n/a 0, /* tp_as_sequence */
3552n/a 0, /* tp_as_mapping */
3553n/a 0, /* tp_hash */
3554n/a 0, /* tp_call */
3555n/a 0, /* tp_str */
3556n/a PyObject_GenericGetAttr, /* tp_getattro */
3557n/a 0, /* tp_setattro */
3558n/a 0, /* tp_as_buffer */
3559n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3560n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
3561n/a accumulate_doc, /* tp_doc */
3562n/a (traverseproc)accumulate_traverse, /* tp_traverse */
3563n/a 0, /* tp_clear */
3564n/a 0, /* tp_richcompare */
3565n/a 0, /* tp_weaklistoffset */
3566n/a PyObject_SelfIter, /* tp_iter */
3567n/a (iternextfunc)accumulate_next, /* tp_iternext */
3568n/a accumulate_methods, /* tp_methods */
3569n/a 0, /* tp_members */
3570n/a 0, /* tp_getset */
3571n/a 0, /* tp_base */
3572n/a 0, /* tp_dict */
3573n/a 0, /* tp_descr_get */
3574n/a 0, /* tp_descr_set */
3575n/a 0, /* tp_dictoffset */
3576n/a 0, /* tp_init */
3577n/a 0, /* tp_alloc */
3578n/a accumulate_new, /* tp_new */
3579n/a PyObject_GC_Del, /* tp_free */
3580n/a};
3581n/a
3582n/a
3583n/a/* compress object ************************************************************/
3584n/a
3585n/a/* Equivalent to:
3586n/a
3587n/a def compress(data, selectors):
3588n/a "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3589n/a return (d for d, s in zip(data, selectors) if s)
3590n/a*/
3591n/a
3592n/atypedef struct {
3593n/a PyObject_HEAD
3594n/a PyObject *data;
3595n/a PyObject *selectors;
3596n/a} compressobject;
3597n/a
3598n/astatic PyTypeObject compress_type;
3599n/a
3600n/astatic PyObject *
3601n/acompress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3602n/a{
3603n/a PyObject *seq1, *seq2;
3604n/a PyObject *data=NULL, *selectors=NULL;
3605n/a compressobject *lz;
3606n/a static char *kwargs[] = {"data", "selectors", NULL};
3607n/a
3608n/a if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3609n/a return NULL;
3610n/a
3611n/a data = PyObject_GetIter(seq1);
3612n/a if (data == NULL)
3613n/a goto fail;
3614n/a selectors = PyObject_GetIter(seq2);
3615n/a if (selectors == NULL)
3616n/a goto fail;
3617n/a
3618n/a /* create compressobject structure */
3619n/a lz = (compressobject *)type->tp_alloc(type, 0);
3620n/a if (lz == NULL)
3621n/a goto fail;
3622n/a lz->data = data;
3623n/a lz->selectors = selectors;
3624n/a return (PyObject *)lz;
3625n/a
3626n/afail:
3627n/a Py_XDECREF(data);
3628n/a Py_XDECREF(selectors);
3629n/a return NULL;
3630n/a}
3631n/a
3632n/astatic void
3633n/acompress_dealloc(compressobject *lz)
3634n/a{
3635n/a PyObject_GC_UnTrack(lz);
3636n/a Py_XDECREF(lz->data);
3637n/a Py_XDECREF(lz->selectors);
3638n/a Py_TYPE(lz)->tp_free(lz);
3639n/a}
3640n/a
3641n/astatic int
3642n/acompress_traverse(compressobject *lz, visitproc visit, void *arg)
3643n/a{
3644n/a Py_VISIT(lz->data);
3645n/a Py_VISIT(lz->selectors);
3646n/a return 0;
3647n/a}
3648n/a
3649n/astatic PyObject *
3650n/acompress_next(compressobject *lz)
3651n/a{
3652n/a PyObject *data = lz->data, *selectors = lz->selectors;
3653n/a PyObject *datum, *selector;
3654n/a PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3655n/a PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3656n/a int ok;
3657n/a
3658n/a while (1) {
3659n/a /* Steps: get datum, get selector, evaluate selector.
3660n/a Order is important (to match the pure python version
3661n/a in terms of which input gets a chance to raise an
3662n/a exception first).
3663n/a */
3664n/a
3665n/a datum = datanext(data);
3666n/a if (datum == NULL)
3667n/a return NULL;
3668n/a
3669n/a selector = selectornext(selectors);
3670n/a if (selector == NULL) {
3671n/a Py_DECREF(datum);
3672n/a return NULL;
3673n/a }
3674n/a
3675n/a ok = PyObject_IsTrue(selector);
3676n/a Py_DECREF(selector);
3677n/a if (ok > 0)
3678n/a return datum;
3679n/a Py_DECREF(datum);
3680n/a if (ok < 0)
3681n/a return NULL;
3682n/a }
3683n/a}
3684n/a
3685n/astatic PyObject *
3686n/acompress_reduce(compressobject *lz)
3687n/a{
3688n/a return Py_BuildValue("O(OO)", Py_TYPE(lz),
3689n/a lz->data, lz->selectors);
3690n/a}
3691n/a
3692n/astatic PyMethodDef compress_methods[] = {
3693n/a {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3694n/a reduce_doc},
3695n/a {NULL, NULL} /* sentinel */
3696n/a};
3697n/a
3698n/aPyDoc_STRVAR(compress_doc,
3699n/a"compress(data, selectors) --> iterator over selected data\n\
3700n/a\n\
3701n/aReturn data elements corresponding to true selector elements.\n\
3702n/aForms a shorter iterator from selected data elements using the\n\
3703n/aselectors to choose the data elements.");
3704n/a
3705n/astatic PyTypeObject compress_type = {
3706n/a PyVarObject_HEAD_INIT(NULL, 0)
3707n/a "itertools.compress", /* tp_name */
3708n/a sizeof(compressobject), /* tp_basicsize */
3709n/a 0, /* tp_itemsize */
3710n/a /* methods */
3711n/a (destructor)compress_dealloc, /* tp_dealloc */
3712n/a 0, /* tp_print */
3713n/a 0, /* tp_getattr */
3714n/a 0, /* tp_setattr */
3715n/a 0, /* tp_reserved */
3716n/a 0, /* tp_repr */
3717n/a 0, /* tp_as_number */
3718n/a 0, /* tp_as_sequence */
3719n/a 0, /* tp_as_mapping */
3720n/a 0, /* tp_hash */
3721n/a 0, /* tp_call */
3722n/a 0, /* tp_str */
3723n/a PyObject_GenericGetAttr, /* tp_getattro */
3724n/a 0, /* tp_setattro */
3725n/a 0, /* tp_as_buffer */
3726n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3727n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
3728n/a compress_doc, /* tp_doc */
3729n/a (traverseproc)compress_traverse, /* tp_traverse */
3730n/a 0, /* tp_clear */
3731n/a 0, /* tp_richcompare */
3732n/a 0, /* tp_weaklistoffset */
3733n/a PyObject_SelfIter, /* tp_iter */
3734n/a (iternextfunc)compress_next, /* tp_iternext */
3735n/a compress_methods, /* tp_methods */
3736n/a 0, /* tp_members */
3737n/a 0, /* tp_getset */
3738n/a 0, /* tp_base */
3739n/a 0, /* tp_dict */
3740n/a 0, /* tp_descr_get */
3741n/a 0, /* tp_descr_set */
3742n/a 0, /* tp_dictoffset */
3743n/a 0, /* tp_init */
3744n/a 0, /* tp_alloc */
3745n/a compress_new, /* tp_new */
3746n/a PyObject_GC_Del, /* tp_free */
3747n/a};
3748n/a
3749n/a
3750n/a/* filterfalse object ************************************************************/
3751n/a
3752n/atypedef struct {
3753n/a PyObject_HEAD
3754n/a PyObject *func;
3755n/a PyObject *it;
3756n/a} filterfalseobject;
3757n/a
3758n/astatic PyTypeObject filterfalse_type;
3759n/a
3760n/astatic PyObject *
3761n/afilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3762n/a{
3763n/a PyObject *func, *seq;
3764n/a PyObject *it;
3765n/a filterfalseobject *lz;
3766n/a
3767n/a if (type == &filterfalse_type &&
3768n/a !_PyArg_NoKeywords("filterfalse()", kwds))
3769n/a return NULL;
3770n/a
3771n/a if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3772n/a return NULL;
3773n/a
3774n/a /* Get iterator. */
3775n/a it = PyObject_GetIter(seq);
3776n/a if (it == NULL)
3777n/a return NULL;
3778n/a
3779n/a /* create filterfalseobject structure */
3780n/a lz = (filterfalseobject *)type->tp_alloc(type, 0);
3781n/a if (lz == NULL) {
3782n/a Py_DECREF(it);
3783n/a return NULL;
3784n/a }
3785n/a Py_INCREF(func);
3786n/a lz->func = func;
3787n/a lz->it = it;
3788n/a
3789n/a return (PyObject *)lz;
3790n/a}
3791n/a
3792n/astatic void
3793n/afilterfalse_dealloc(filterfalseobject *lz)
3794n/a{
3795n/a PyObject_GC_UnTrack(lz);
3796n/a Py_XDECREF(lz->func);
3797n/a Py_XDECREF(lz->it);
3798n/a Py_TYPE(lz)->tp_free(lz);
3799n/a}
3800n/a
3801n/astatic int
3802n/afilterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
3803n/a{
3804n/a Py_VISIT(lz->it);
3805n/a Py_VISIT(lz->func);
3806n/a return 0;
3807n/a}
3808n/a
3809n/astatic PyObject *
3810n/afilterfalse_next(filterfalseobject *lz)
3811n/a{
3812n/a PyObject *item;
3813n/a PyObject *it = lz->it;
3814n/a long ok;
3815n/a PyObject *(*iternext)(PyObject *);
3816n/a
3817n/a iternext = *Py_TYPE(it)->tp_iternext;
3818n/a for (;;) {
3819n/a item = iternext(it);
3820n/a if (item == NULL)
3821n/a return NULL;
3822n/a
3823n/a if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3824n/a ok = PyObject_IsTrue(item);
3825n/a } else {
3826n/a PyObject *good;
3827n/a good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
3828n/a if (good == NULL) {
3829n/a Py_DECREF(item);
3830n/a return NULL;
3831n/a }
3832n/a ok = PyObject_IsTrue(good);
3833n/a Py_DECREF(good);
3834n/a }
3835n/a if (ok == 0)
3836n/a return item;
3837n/a Py_DECREF(item);
3838n/a if (ok < 0)
3839n/a return NULL;
3840n/a }
3841n/a}
3842n/a
3843n/astatic PyObject *
3844n/afilterfalse_reduce(filterfalseobject *lz)
3845n/a{
3846n/a return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3847n/a}
3848n/a
3849n/astatic PyMethodDef filterfalse_methods[] = {
3850n/a {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3851n/a reduce_doc},
3852n/a {NULL, NULL} /* sentinel */
3853n/a};
3854n/a
3855n/aPyDoc_STRVAR(filterfalse_doc,
3856n/a"filterfalse(function or None, sequence) --> filterfalse object\n\
3857n/a\n\
3858n/aReturn those items of sequence for which function(item) is false.\n\
3859n/aIf function is None, return the items that are false.");
3860n/a
3861n/astatic PyTypeObject filterfalse_type = {
3862n/a PyVarObject_HEAD_INIT(NULL, 0)
3863n/a "itertools.filterfalse", /* tp_name */
3864n/a sizeof(filterfalseobject), /* tp_basicsize */
3865n/a 0, /* tp_itemsize */
3866n/a /* methods */
3867n/a (destructor)filterfalse_dealloc, /* tp_dealloc */
3868n/a 0, /* tp_print */
3869n/a 0, /* tp_getattr */
3870n/a 0, /* tp_setattr */
3871n/a 0, /* tp_reserved */
3872n/a 0, /* tp_repr */
3873n/a 0, /* tp_as_number */
3874n/a 0, /* tp_as_sequence */
3875n/a 0, /* tp_as_mapping */
3876n/a 0, /* tp_hash */
3877n/a 0, /* tp_call */
3878n/a 0, /* tp_str */
3879n/a PyObject_GenericGetAttr, /* tp_getattro */
3880n/a 0, /* tp_setattro */
3881n/a 0, /* tp_as_buffer */
3882n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3883n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
3884n/a filterfalse_doc, /* tp_doc */
3885n/a (traverseproc)filterfalse_traverse, /* tp_traverse */
3886n/a 0, /* tp_clear */
3887n/a 0, /* tp_richcompare */
3888n/a 0, /* tp_weaklistoffset */
3889n/a PyObject_SelfIter, /* tp_iter */
3890n/a (iternextfunc)filterfalse_next, /* tp_iternext */
3891n/a filterfalse_methods, /* tp_methods */
3892n/a 0, /* tp_members */
3893n/a 0, /* tp_getset */
3894n/a 0, /* tp_base */
3895n/a 0, /* tp_dict */
3896n/a 0, /* tp_descr_get */
3897n/a 0, /* tp_descr_set */
3898n/a 0, /* tp_dictoffset */
3899n/a 0, /* tp_init */
3900n/a 0, /* tp_alloc */
3901n/a filterfalse_new, /* tp_new */
3902n/a PyObject_GC_Del, /* tp_free */
3903n/a};
3904n/a
3905n/a
3906n/a/* count object ************************************************************/
3907n/a
3908n/atypedef struct {
3909n/a PyObject_HEAD
3910n/a Py_ssize_t cnt;
3911n/a PyObject *long_cnt;
3912n/a PyObject *long_step;
3913n/a} countobject;
3914n/a
3915n/a/* Counting logic and invariants:
3916n/a
3917n/afast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3918n/a
3919n/a assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
3920n/a Advances with: cnt += 1
3921n/a When count hits Y_SSIZE_T_MAX, switch to slow_mode.
3922n/a
3923n/aslow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3924n/a
3925n/a assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3926n/a All counting is done with python objects (no overflows or underflows).
3927n/a Advances with: long_cnt += long_step
3928n/a Step may be zero -- effectively a slow version of repeat(cnt).
3929n/a Either long_cnt or long_step may be a float, Fraction, or Decimal.
3930n/a*/
3931n/a
3932n/astatic PyTypeObject count_type;
3933n/a
3934n/astatic PyObject *
3935n/acount_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3936n/a{
3937n/a countobject *lz;
3938n/a int fast_mode;
3939n/a Py_ssize_t cnt = 0;
3940n/a PyObject *long_cnt = NULL;
3941n/a PyObject *long_step = NULL;
3942n/a long step;
3943n/a static char *kwlist[] = {"start", "step", 0};
3944n/a
3945n/a if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3946n/a kwlist, &long_cnt, &long_step))
3947n/a return NULL;
3948n/a
3949n/a if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3950n/a (long_step != NULL && !PyNumber_Check(long_step))) {
3951n/a PyErr_SetString(PyExc_TypeError, "a number is required");
3952n/a return NULL;
3953n/a }
3954n/a
3955n/a fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3956n/a (long_step == NULL || PyLong_Check(long_step));
3957n/a
3958n/a /* If not specified, start defaults to 0 */
3959n/a if (long_cnt != NULL) {
3960n/a if (fast_mode) {
3961n/a assert(PyLong_Check(long_cnt));
3962n/a cnt = PyLong_AsSsize_t(long_cnt);
3963n/a if (cnt == -1 && PyErr_Occurred()) {
3964n/a PyErr_Clear();
3965n/a fast_mode = 0;
3966n/a }
3967n/a }
3968n/a Py_INCREF(long_cnt);
3969n/a } else {
3970n/a cnt = 0;
3971n/a long_cnt = PyLong_FromLong(0);
3972n/a if (long_cnt == NULL) {
3973n/a return NULL;
3974n/a }
3975n/a }
3976n/a
3977n/a /* If not specified, step defaults to 1 */
3978n/a if (long_step == NULL) {
3979n/a long_step = PyLong_FromLong(1);
3980n/a if (long_step == NULL) {
3981n/a Py_DECREF(long_cnt);
3982n/a return NULL;
3983n/a }
3984n/a } else
3985n/a Py_INCREF(long_step);
3986n/a
3987n/a assert(long_cnt != NULL && long_step != NULL);
3988n/a
3989n/a /* Fast mode only works when the step is 1 */
3990n/a if (fast_mode) {
3991n/a assert(PyLong_Check(long_step));
3992n/a step = PyLong_AsLong(long_step);
3993n/a if (step != 1) {
3994n/a fast_mode = 0;
3995n/a if (step == -1 && PyErr_Occurred())
3996n/a PyErr_Clear();
3997n/a }
3998n/a }
3999n/a
4000n/a if (fast_mode)
4001n/a Py_CLEAR(long_cnt);
4002n/a else
4003n/a cnt = PY_SSIZE_T_MAX;
4004n/a
4005n/a assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4006n/a (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4007n/a assert(!fast_mode ||
4008n/a (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
4009n/a
4010n/a /* create countobject structure */
4011n/a lz = (countobject *)type->tp_alloc(type, 0);
4012n/a if (lz == NULL) {
4013n/a Py_XDECREF(long_cnt);
4014n/a return NULL;
4015n/a }
4016n/a lz->cnt = cnt;
4017n/a lz->long_cnt = long_cnt;
4018n/a lz->long_step = long_step;
4019n/a
4020n/a return (PyObject *)lz;
4021n/a}
4022n/a
4023n/astatic void
4024n/acount_dealloc(countobject *lz)
4025n/a{
4026n/a PyObject_GC_UnTrack(lz);
4027n/a Py_XDECREF(lz->long_cnt);
4028n/a Py_XDECREF(lz->long_step);
4029n/a Py_TYPE(lz)->tp_free(lz);
4030n/a}
4031n/a
4032n/astatic int
4033n/acount_traverse(countobject *lz, visitproc visit, void *arg)
4034n/a{
4035n/a Py_VISIT(lz->long_cnt);
4036n/a Py_VISIT(lz->long_step);
4037n/a return 0;
4038n/a}
4039n/a
4040n/astatic PyObject *
4041n/acount_nextlong(countobject *lz)
4042n/a{
4043n/a PyObject *long_cnt;
4044n/a PyObject *stepped_up;
4045n/a
4046n/a long_cnt = lz->long_cnt;
4047n/a if (long_cnt == NULL) {
4048n/a /* Switch to slow_mode */
4049n/a long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4050n/a if (long_cnt == NULL)
4051n/a return NULL;
4052n/a }
4053n/a assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
4054n/a
4055n/a stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4056n/a if (stepped_up == NULL)
4057n/a return NULL;
4058n/a lz->long_cnt = stepped_up;
4059n/a return long_cnt;
4060n/a}
4061n/a
4062n/astatic PyObject *
4063n/acount_next(countobject *lz)
4064n/a{
4065n/a if (lz->cnt == PY_SSIZE_T_MAX)
4066n/a return count_nextlong(lz);
4067n/a return PyLong_FromSsize_t(lz->cnt++);
4068n/a}
4069n/a
4070n/astatic PyObject *
4071n/acount_repr(countobject *lz)
4072n/a{
4073n/a if (lz->cnt != PY_SSIZE_T_MAX)
4074n/a return PyUnicode_FromFormat("count(%zd)", lz->cnt);
4075n/a
4076n/a if (PyLong_Check(lz->long_step)) {
4077n/a long step = PyLong_AsLong(lz->long_step);
4078n/a if (step == -1 && PyErr_Occurred()) {
4079n/a PyErr_Clear();
4080n/a }
4081n/a if (step == 1) {
4082n/a /* Don't display step when it is an integer equal to 1 */
4083n/a return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4084n/a }
4085n/a }
4086n/a return PyUnicode_FromFormat("count(%R, %R)",
4087n/a lz->long_cnt, lz->long_step);
4088n/a}
4089n/a
4090n/astatic PyObject *
4091n/acount_reduce(countobject *lz)
4092n/a{
4093n/a if (lz->cnt == PY_SSIZE_T_MAX)
4094n/a return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4095n/a return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
4096n/a}
4097n/a
4098n/astatic PyMethodDef count_methods[] = {
4099n/a {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
4100n/a reduce_doc},
4101n/a {NULL, NULL} /* sentinel */
4102n/a};
4103n/a
4104n/aPyDoc_STRVAR(count_doc,
4105n/a "count(start=0, step=1) --> count object\n\
4106n/a\n\
4107n/aReturn a count object whose .__next__() method returns consecutive values.\n\
4108n/aEquivalent to:\n\n\
4109n/a def count(firstval=0, step=1):\n\
4110n/a x = firstval\n\
4111n/a while 1:\n\
4112n/a yield x\n\
4113n/a x += step\n");
4114n/a
4115n/astatic PyTypeObject count_type = {
4116n/a PyVarObject_HEAD_INIT(NULL, 0)
4117n/a "itertools.count", /* tp_name */
4118n/a sizeof(countobject), /* tp_basicsize */
4119n/a 0, /* tp_itemsize */
4120n/a /* methods */
4121n/a (destructor)count_dealloc, /* tp_dealloc */
4122n/a 0, /* tp_print */
4123n/a 0, /* tp_getattr */
4124n/a 0, /* tp_setattr */
4125n/a 0, /* tp_reserved */
4126n/a (reprfunc)count_repr, /* tp_repr */
4127n/a 0, /* tp_as_number */
4128n/a 0, /* tp_as_sequence */
4129n/a 0, /* tp_as_mapping */
4130n/a 0, /* tp_hash */
4131n/a 0, /* tp_call */
4132n/a 0, /* tp_str */
4133n/a PyObject_GenericGetAttr, /* tp_getattro */
4134n/a 0, /* tp_setattro */
4135n/a 0, /* tp_as_buffer */
4136n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4137n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
4138n/a count_doc, /* tp_doc */
4139n/a (traverseproc)count_traverse, /* tp_traverse */
4140n/a 0, /* tp_clear */
4141n/a 0, /* tp_richcompare */
4142n/a 0, /* tp_weaklistoffset */
4143n/a PyObject_SelfIter, /* tp_iter */
4144n/a (iternextfunc)count_next, /* tp_iternext */
4145n/a count_methods, /* tp_methods */
4146n/a 0, /* tp_members */
4147n/a 0, /* tp_getset */
4148n/a 0, /* tp_base */
4149n/a 0, /* tp_dict */
4150n/a 0, /* tp_descr_get */
4151n/a 0, /* tp_descr_set */
4152n/a 0, /* tp_dictoffset */
4153n/a 0, /* tp_init */
4154n/a 0, /* tp_alloc */
4155n/a count_new, /* tp_new */
4156n/a PyObject_GC_Del, /* tp_free */
4157n/a};
4158n/a
4159n/a
4160n/a/* repeat object ************************************************************/
4161n/a
4162n/atypedef struct {
4163n/a PyObject_HEAD
4164n/a PyObject *element;
4165n/a Py_ssize_t cnt;
4166n/a} repeatobject;
4167n/a
4168n/astatic PyTypeObject repeat_type;
4169n/a
4170n/astatic PyObject *
4171n/arepeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4172n/a{
4173n/a repeatobject *ro;
4174n/a PyObject *element;
4175n/a Py_ssize_t cnt = -1, n_kwds = 0;
4176n/a static char *kwargs[] = {"object", "times", NULL};
4177n/a
4178n/a if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4179n/a &element, &cnt))
4180n/a return NULL;
4181n/a
4182n/a if (kwds != NULL)
4183n/a n_kwds = PyDict_GET_SIZE(kwds);
4184n/a /* Does user supply times argument? */
4185n/a if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
4186n/a cnt = 0;
4187n/a
4188n/a ro = (repeatobject *)type->tp_alloc(type, 0);
4189n/a if (ro == NULL)
4190n/a return NULL;
4191n/a Py_INCREF(element);
4192n/a ro->element = element;
4193n/a ro->cnt = cnt;
4194n/a return (PyObject *)ro;
4195n/a}
4196n/a
4197n/astatic void
4198n/arepeat_dealloc(repeatobject *ro)
4199n/a{
4200n/a PyObject_GC_UnTrack(ro);
4201n/a Py_XDECREF(ro->element);
4202n/a Py_TYPE(ro)->tp_free(ro);
4203n/a}
4204n/a
4205n/astatic int
4206n/arepeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4207n/a{
4208n/a Py_VISIT(ro->element);
4209n/a return 0;
4210n/a}
4211n/a
4212n/astatic PyObject *
4213n/arepeat_next(repeatobject *ro)
4214n/a{
4215n/a if (ro->cnt == 0)
4216n/a return NULL;
4217n/a if (ro->cnt > 0)
4218n/a ro->cnt--;
4219n/a Py_INCREF(ro->element);
4220n/a return ro->element;
4221n/a}
4222n/a
4223n/astatic PyObject *
4224n/arepeat_repr(repeatobject *ro)
4225n/a{
4226n/a if (ro->cnt == -1)
4227n/a return PyUnicode_FromFormat("repeat(%R)", ro->element);
4228n/a else
4229n/a return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4230n/a}
4231n/a
4232n/astatic PyObject *
4233n/arepeat_len(repeatobject *ro)
4234n/a{
4235n/a if (ro->cnt == -1) {
4236n/a PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4237n/a return NULL;
4238n/a }
4239n/a return PyLong_FromSize_t(ro->cnt);
4240n/a}
4241n/a
4242n/aPyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
4243n/a
4244n/astatic PyObject *
4245n/arepeat_reduce(repeatobject *ro)
4246n/a{
4247n/a /* unpickle this so that a new repeat iterator is constructed with an
4248n/a * object, then call __setstate__ on it to set cnt
4249n/a */
4250n/a if (ro->cnt >= 0)
4251n/a return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4252n/a else
4253n/a return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4254n/a}
4255n/a
4256n/astatic PyMethodDef repeat_methods[] = {
4257n/a {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
4258n/a {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
4259n/a {NULL, NULL} /* sentinel */
4260n/a};
4261n/a
4262n/aPyDoc_STRVAR(repeat_doc,
4263n/a"repeat(object [,times]) -> create an iterator which returns the object\n\
4264n/afor the specified number of times. If not specified, returns the object\n\
4265n/aendlessly.");
4266n/a
4267n/astatic PyTypeObject repeat_type = {
4268n/a PyVarObject_HEAD_INIT(NULL, 0)
4269n/a "itertools.repeat", /* tp_name */
4270n/a sizeof(repeatobject), /* tp_basicsize */
4271n/a 0, /* tp_itemsize */
4272n/a /* methods */
4273n/a (destructor)repeat_dealloc, /* tp_dealloc */
4274n/a 0, /* tp_print */
4275n/a 0, /* tp_getattr */
4276n/a 0, /* tp_setattr */
4277n/a 0, /* tp_reserved */
4278n/a (reprfunc)repeat_repr, /* tp_repr */
4279n/a 0, /* tp_as_number */
4280n/a 0, /* tp_as_sequence */
4281n/a 0, /* tp_as_mapping */
4282n/a 0, /* tp_hash */
4283n/a 0, /* tp_call */
4284n/a 0, /* tp_str */
4285n/a PyObject_GenericGetAttr, /* tp_getattro */
4286n/a 0, /* tp_setattro */
4287n/a 0, /* tp_as_buffer */
4288n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4289n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
4290n/a repeat_doc, /* tp_doc */
4291n/a (traverseproc)repeat_traverse, /* tp_traverse */
4292n/a 0, /* tp_clear */
4293n/a 0, /* tp_richcompare */
4294n/a 0, /* tp_weaklistoffset */
4295n/a PyObject_SelfIter, /* tp_iter */
4296n/a (iternextfunc)repeat_next, /* tp_iternext */
4297n/a repeat_methods, /* tp_methods */
4298n/a 0, /* tp_members */
4299n/a 0, /* tp_getset */
4300n/a 0, /* tp_base */
4301n/a 0, /* tp_dict */
4302n/a 0, /* tp_descr_get */
4303n/a 0, /* tp_descr_set */
4304n/a 0, /* tp_dictoffset */
4305n/a 0, /* tp_init */
4306n/a 0, /* tp_alloc */
4307n/a repeat_new, /* tp_new */
4308n/a PyObject_GC_Del, /* tp_free */
4309n/a};
4310n/a
4311n/a/* ziplongest object *********************************************************/
4312n/a
4313n/atypedef struct {
4314n/a PyObject_HEAD
4315n/a Py_ssize_t tuplesize;
4316n/a Py_ssize_t numactive;
4317n/a PyObject *ittuple; /* tuple of iterators */
4318n/a PyObject *result;
4319n/a PyObject *fillvalue;
4320n/a} ziplongestobject;
4321n/a
4322n/astatic PyTypeObject ziplongest_type;
4323n/a
4324n/astatic PyObject *
4325n/azip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4326n/a{
4327n/a ziplongestobject *lz;
4328n/a Py_ssize_t i;
4329n/a PyObject *ittuple; /* tuple of iterators */
4330n/a PyObject *result;
4331n/a PyObject *fillvalue = Py_None;
4332n/a Py_ssize_t tuplesize = PySequence_Length(args);
4333n/a
4334n/a if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
4335n/a fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4336n/a if (fillvalue == NULL || PyDict_GET_SIZE(kwds) > 1) {
4337n/a PyErr_SetString(PyExc_TypeError,
4338n/a "zip_longest() got an unexpected keyword argument");
4339n/a return NULL;
4340n/a }
4341n/a }
4342n/a
4343n/a /* args must be a tuple */
4344n/a assert(PyTuple_Check(args));
4345n/a
4346n/a /* obtain iterators */
4347n/a ittuple = PyTuple_New(tuplesize);
4348n/a if (ittuple == NULL)
4349n/a return NULL;
4350n/a for (i=0; i < tuplesize; i++) {
4351n/a PyObject *item = PyTuple_GET_ITEM(args, i);
4352n/a PyObject *it = PyObject_GetIter(item);
4353n/a if (it == NULL) {
4354n/a if (PyErr_ExceptionMatches(PyExc_TypeError))
4355n/a PyErr_Format(PyExc_TypeError,
4356n/a "zip_longest argument #%zd must support iteration",
4357n/a i+1);
4358n/a Py_DECREF(ittuple);
4359n/a return NULL;
4360n/a }
4361n/a PyTuple_SET_ITEM(ittuple, i, it);
4362n/a }
4363n/a
4364n/a /* create a result holder */
4365n/a result = PyTuple_New(tuplesize);
4366n/a if (result == NULL) {
4367n/a Py_DECREF(ittuple);
4368n/a return NULL;
4369n/a }
4370n/a for (i=0 ; i < tuplesize ; i++) {
4371n/a Py_INCREF(Py_None);
4372n/a PyTuple_SET_ITEM(result, i, Py_None);
4373n/a }
4374n/a
4375n/a /* create ziplongestobject structure */
4376n/a lz = (ziplongestobject *)type->tp_alloc(type, 0);
4377n/a if (lz == NULL) {
4378n/a Py_DECREF(ittuple);
4379n/a Py_DECREF(result);
4380n/a return NULL;
4381n/a }
4382n/a lz->ittuple = ittuple;
4383n/a lz->tuplesize = tuplesize;
4384n/a lz->numactive = tuplesize;
4385n/a lz->result = result;
4386n/a Py_INCREF(fillvalue);
4387n/a lz->fillvalue = fillvalue;
4388n/a return (PyObject *)lz;
4389n/a}
4390n/a
4391n/astatic void
4392n/azip_longest_dealloc(ziplongestobject *lz)
4393n/a{
4394n/a PyObject_GC_UnTrack(lz);
4395n/a Py_XDECREF(lz->ittuple);
4396n/a Py_XDECREF(lz->result);
4397n/a Py_XDECREF(lz->fillvalue);
4398n/a Py_TYPE(lz)->tp_free(lz);
4399n/a}
4400n/a
4401n/astatic int
4402n/azip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
4403n/a{
4404n/a Py_VISIT(lz->ittuple);
4405n/a Py_VISIT(lz->result);
4406n/a Py_VISIT(lz->fillvalue);
4407n/a return 0;
4408n/a}
4409n/a
4410n/astatic PyObject *
4411n/azip_longest_next(ziplongestobject *lz)
4412n/a{
4413n/a Py_ssize_t i;
4414n/a Py_ssize_t tuplesize = lz->tuplesize;
4415n/a PyObject *result = lz->result;
4416n/a PyObject *it;
4417n/a PyObject *item;
4418n/a PyObject *olditem;
4419n/a
4420n/a if (tuplesize == 0)
4421n/a return NULL;
4422n/a if (lz->numactive == 0)
4423n/a return NULL;
4424n/a if (Py_REFCNT(result) == 1) {
4425n/a Py_INCREF(result);
4426n/a for (i=0 ; i < tuplesize ; i++) {
4427n/a it = PyTuple_GET_ITEM(lz->ittuple, i);
4428n/a if (it == NULL) {
4429n/a Py_INCREF(lz->fillvalue);
4430n/a item = lz->fillvalue;
4431n/a } else {
4432n/a item = PyIter_Next(it);
4433n/a if (item == NULL) {
4434n/a lz->numactive -= 1;
4435n/a if (lz->numactive == 0 || PyErr_Occurred()) {
4436n/a lz->numactive = 0;
4437n/a Py_DECREF(result);
4438n/a return NULL;
4439n/a } else {
4440n/a Py_INCREF(lz->fillvalue);
4441n/a item = lz->fillvalue;
4442n/a PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4443n/a Py_DECREF(it);
4444n/a }
4445n/a }
4446n/a }
4447n/a olditem = PyTuple_GET_ITEM(result, i);
4448n/a PyTuple_SET_ITEM(result, i, item);
4449n/a Py_DECREF(olditem);
4450n/a }
4451n/a } else {
4452n/a result = PyTuple_New(tuplesize);
4453n/a if (result == NULL)
4454n/a return NULL;
4455n/a for (i=0 ; i < tuplesize ; i++) {
4456n/a it = PyTuple_GET_ITEM(lz->ittuple, i);
4457n/a if (it == NULL) {
4458n/a Py_INCREF(lz->fillvalue);
4459n/a item = lz->fillvalue;
4460n/a } else {
4461n/a item = PyIter_Next(it);
4462n/a if (item == NULL) {
4463n/a lz->numactive -= 1;
4464n/a if (lz->numactive == 0 || PyErr_Occurred()) {
4465n/a lz->numactive = 0;
4466n/a Py_DECREF(result);
4467n/a return NULL;
4468n/a } else {
4469n/a Py_INCREF(lz->fillvalue);
4470n/a item = lz->fillvalue;
4471n/a PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4472n/a Py_DECREF(it);
4473n/a }
4474n/a }
4475n/a }
4476n/a PyTuple_SET_ITEM(result, i, item);
4477n/a }
4478n/a }
4479n/a return result;
4480n/a}
4481n/a
4482n/astatic PyObject *
4483n/azip_longest_reduce(ziplongestobject *lz)
4484n/a{
4485n/a
4486n/a /* Create a new tuple with empty sequences where appropriate to pickle.
4487n/a * Then use setstate to set the fillvalue
4488n/a */
4489n/a int i;
4490n/a PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4491n/a
4492n/a if (args == NULL)
4493n/a return NULL;
4494n/a for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4495n/a PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4496n/a if (elem == NULL) {
4497n/a elem = PyTuple_New(0);
4498n/a if (elem == NULL) {
4499n/a Py_DECREF(args);
4500n/a return NULL;
4501n/a }
4502n/a } else
4503n/a Py_INCREF(elem);
4504n/a PyTuple_SET_ITEM(args, i, elem);
4505n/a }
4506n/a return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4507n/a}
4508n/a
4509n/astatic PyObject *
4510n/azip_longest_setstate(ziplongestobject *lz, PyObject *state)
4511n/a{
4512n/a Py_INCREF(state);
4513n/a Py_XSETREF(lz->fillvalue, state);
4514n/a Py_RETURN_NONE;
4515n/a}
4516n/a
4517n/astatic PyMethodDef zip_longest_methods[] = {
4518n/a {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4519n/a reduce_doc},
4520n/a {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4521n/a setstate_doc},
4522n/a {NULL, NULL} /* sentinel */
4523n/a};
4524n/a
4525n/aPyDoc_STRVAR(zip_longest_doc,
4526n/a"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
4527n/a\n\
4528n/aReturn a zip_longest object whose .__next__() method returns a tuple where\n\
4529n/athe i-th element comes from the i-th iterable argument. The .__next__()\n\
4530n/amethod continues until the longest iterable in the argument sequence\n\
4531n/ais exhausted and then it raises StopIteration. When the shorter iterables\n\
4532n/aare exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4533n/adefaults to None or can be specified by a keyword argument.\n\
4534n/a");
4535n/a
4536n/astatic PyTypeObject ziplongest_type = {
4537n/a PyVarObject_HEAD_INIT(NULL, 0)
4538n/a "itertools.zip_longest", /* tp_name */
4539n/a sizeof(ziplongestobject), /* tp_basicsize */
4540n/a 0, /* tp_itemsize */
4541n/a /* methods */
4542n/a (destructor)zip_longest_dealloc, /* tp_dealloc */
4543n/a 0, /* tp_print */
4544n/a 0, /* tp_getattr */
4545n/a 0, /* tp_setattr */
4546n/a 0, /* tp_reserved */
4547n/a 0, /* tp_repr */
4548n/a 0, /* tp_as_number */
4549n/a 0, /* tp_as_sequence */
4550n/a 0, /* tp_as_mapping */
4551n/a 0, /* tp_hash */
4552n/a 0, /* tp_call */
4553n/a 0, /* tp_str */
4554n/a PyObject_GenericGetAttr, /* tp_getattro */
4555n/a 0, /* tp_setattro */
4556n/a 0, /* tp_as_buffer */
4557n/a Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4558n/a Py_TPFLAGS_BASETYPE, /* tp_flags */
4559n/a zip_longest_doc, /* tp_doc */
4560n/a (traverseproc)zip_longest_traverse, /* tp_traverse */
4561n/a 0, /* tp_clear */
4562n/a 0, /* tp_richcompare */
4563n/a 0, /* tp_weaklistoffset */
4564n/a PyObject_SelfIter, /* tp_iter */
4565n/a (iternextfunc)zip_longest_next, /* tp_iternext */
4566n/a zip_longest_methods, /* tp_methods */
4567n/a 0, /* tp_members */
4568n/a 0, /* tp_getset */
4569n/a 0, /* tp_base */
4570n/a 0, /* tp_dict */
4571n/a 0, /* tp_descr_get */
4572n/a 0, /* tp_descr_set */
4573n/a 0, /* tp_dictoffset */
4574n/a 0, /* tp_init */
4575n/a 0, /* tp_alloc */
4576n/a zip_longest_new, /* tp_new */
4577n/a PyObject_GC_Del, /* tp_free */
4578n/a};
4579n/a
4580n/a/* module level code ********************************************************/
4581n/a
4582n/aPyDoc_STRVAR(module_doc,
4583n/a"Functional tools for creating and using iterators.\n\
4584n/a\n\
4585n/aInfinite iterators:\n\
4586n/acount(start=0, step=1) --> start, start+step, start+2*step, ...\n\
4587n/acycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
4588n/arepeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
4589n/a\n\
4590n/aIterators terminating on the shortest input sequence:\n\
4591n/aaccumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
4592n/achain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
4593n/achain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
4594n/acompress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4595n/adropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4596n/agroupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
4597n/afilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
4598n/aislice(seq, [start,] stop [, step]) --> elements from\n\
4599n/a seq[start:stop:step]\n\
4600n/astarmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
4601n/atee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
4602n/atakewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
4603n/azip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4604n/a\n\
4605n/aCombinatoric generators:\n\
4606n/aproduct(p, q, ... [repeat=1]) --> cartesian product\n\
4607n/apermutations(p[, r])\n\
4608n/acombinations(p, r)\n\
4609n/acombinations_with_replacement(p, r)\n\
4610n/a");
4611n/a
4612n/a
4613n/astatic PyMethodDef module_methods[] = {
4614n/a {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4615n/a {NULL, NULL} /* sentinel */
4616n/a};
4617n/a
4618n/a
4619n/astatic struct PyModuleDef itertoolsmodule = {
4620n/a PyModuleDef_HEAD_INIT,
4621n/a "itertools",
4622n/a module_doc,
4623n/a -1,
4624n/a module_methods,
4625n/a NULL,
4626n/a NULL,
4627n/a NULL,
4628n/a NULL
4629n/a};
4630n/a
4631n/aPyMODINIT_FUNC
4632n/aPyInit_itertools(void)
4633n/a{
4634n/a int i;
4635n/a PyObject *m;
4636n/a char *name;
4637n/a PyTypeObject *typelist[] = {
4638n/a &accumulate_type,
4639n/a &combinations_type,
4640n/a &cwr_type,
4641n/a &cycle_type,
4642n/a &dropwhile_type,
4643n/a &takewhile_type,
4644n/a &islice_type,
4645n/a &starmap_type,
4646n/a &chain_type,
4647n/a &compress_type,
4648n/a &filterfalse_type,
4649n/a &count_type,
4650n/a &ziplongest_type,
4651n/a &permutations_type,
4652n/a &product_type,
4653n/a &repeat_type,
4654n/a &groupby_type,
4655n/a &_grouper_type,
4656n/a &tee_type,
4657n/a &teedataobject_type,
4658n/a NULL
4659n/a };
4660n/a
4661n/a Py_TYPE(&teedataobject_type) = &PyType_Type;
4662n/a m = PyModule_Create(&itertoolsmodule);
4663n/a if (m == NULL)
4664n/a return NULL;
4665n/a
4666n/a for (i=0 ; typelist[i] != NULL ; i++) {
4667n/a if (PyType_Ready(typelist[i]) < 0)
4668n/a return NULL;
4669n/a name = strchr(typelist[i]->tp_name, '.');
4670n/a assert (name != NULL);
4671n/a Py_INCREF(typelist[i]);
4672n/a PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4673n/a }
4674n/a
4675n/a return m;
4676n/a}