Coverage Report

Created: 2026-02-26 06:25

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