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