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