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