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