Coverage Report

Created: 2026-02-26 06:53

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/src/cpython/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
15.5k
{
45
15.5k
    void *state = _PyModule_GetState(mod);
46
15.5k
    assert(state != NULL);
47
15.5k
    return (itertools_state *)state;
48
15.5k
}
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
14.0k
{
63
14.0k
    PyObject *mod = PyType_GetModuleByDef(tp, &itertoolsmodule);
64
14.0k
    assert(mod != NULL);
65
14.0k
    return get_module_state(mod);
66
14.0k
}
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
46
#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
14
#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
2
{
299
2
    PyObject *it;
300
2
    pairwiseobject *po;
301
302
2
    it = PyObject_GetIter(iterable);
303
2
    if (it == NULL) {
304
0
        return NULL;
305
0
    }
306
2
    po = (pairwiseobject *)type->tp_alloc(type, 0);
307
2
    if (po == NULL) {
308
0
        Py_DECREF(it);
309
0
        return NULL;
310
0
    }
311
2
    po->it = it;
312
2
    po->old = NULL;
313
2
    po->result = PyTuple_Pack(2, Py_None, Py_None);
314
2
    if (po->result == NULL) {
315
0
        Py_DECREF(po);
316
0
        return NULL;
317
0
    }
318
2
    return (PyObject *)po;
319
2
}
320
321
static void
322
pairwise_dealloc(PyObject *op)
323
2
{
324
2
    pairwiseobject *po = pairwiseobject_CAST(op);
325
2
    PyTypeObject *tp = Py_TYPE(po);
326
2
    PyObject_GC_UnTrack(po);
327
2
    Py_XDECREF(po->it);
328
2
    Py_XDECREF(po->old);
329
2
    Py_XDECREF(po->result);
330
2
    tp->tp_free(po);
331
2
    Py_DECREF(tp);
332
2
}
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
12
{
348
12
    pairwiseobject *po = pairwiseobject_CAST(op);
349
12
    PyObject *it = po->it;
350
12
    PyObject *old = po->old;
351
12
    PyObject *new, *result;
352
353
12
    if (it == NULL) {
354
0
        return NULL;
355
0
    }
356
12
    if (old == NULL) {
357
2
        old = (*Py_TYPE(it)->tp_iternext)(it);
358
2
        Py_XSETREF(po->old, old);
359
2
        if (old == NULL) {
360
0
            Py_CLEAR(po->it);
361
0
            return NULL;
362
0
        }
363
2
        it = po->it;
364
2
        if (it == NULL) {
365
0
            Py_CLEAR(po->old);
366
0
            return NULL;
367
0
        }
368
2
    }
369
12
    Py_INCREF(old);
370
12
    new = (*Py_TYPE(it)->tp_iternext)(it);
371
12
    if (new == NULL) {
372
2
        Py_CLEAR(po->it);
373
2
        Py_CLEAR(po->old);
374
2
        Py_DECREF(old);
375
2
        return NULL;
376
2
    }
377
378
10
    result = po->result;
379
10
    if (_PyObject_IsUniquelyReferenced(result)) {
380
10
        Py_INCREF(result);
381
10
        PyObject *last_old = PyTuple_GET_ITEM(result, 0);
382
10
        PyObject *last_new = PyTuple_GET_ITEM(result, 1);
383
10
        PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
384
10
        PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
385
10
        Py_DECREF(last_old);
386
10
        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
10
        _PyTuple_Recycle(result);
390
10
    }
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
10
    Py_XSETREF(po->old, new);
400
10
    Py_DECREF(old);
401
10
    return result;
402
12
}
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
127k
#define chainobject_CAST(op)    ((chainobject *)(op))
1814
1815
static PyObject *
1816
chain_new_internal(PyTypeObject *type, PyObject *source)
1817
14.0k
{
1818
14.0k
    chainobject *lz;
1819
1820
14.0k
    lz = (chainobject *)type->tp_alloc(type, 0);
1821
14.0k
    if (lz == NULL) {
1822
0
        Py_DECREF(source);
1823
0
        return NULL;
1824
0
    }
1825
1826
14.0k
    lz->source = source;
1827
14.0k
    lz->active = NULL;
1828
14.0k
    return (PyObject *)lz;
1829
14.0k
}
1830
1831
static PyObject *
1832
chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1833
13.9k
{
1834
13.9k
    PyObject *source;
1835
1836
13.9k
    itertools_state *state = find_state_by_type(type);
1837
13.9k
    PyTypeObject *chain_type = state->chain_type;
1838
13.9k
    if ((type == chain_type || type->tp_init == chain_type->tp_init) &&
1839
13.9k
        !_PyArg_NoKeywords("chain", kwds))
1840
0
        return NULL;
1841
1842
13.9k
    source = PyObject_GetIter(args);
1843
13.9k
    if (source == NULL)
1844
0
        return NULL;
1845
1846
13.9k
    return chain_new_internal(type, source);
1847
13.9k
}
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
82
{
1862
82
    PyObject *source;
1863
1864
82
    source = PyObject_GetIter(arg);
1865
82
    if (source == NULL)
1866
0
        return NULL;
1867
1868
82
    return chain_new_internal(type, source);
1869
82
}
1870
1871
static void
1872
chain_dealloc(PyObject *op)
1873
14.0k
{
1874
14.0k
    chainobject *lz = chainobject_CAST(op);
1875
14.0k
    PyTypeObject *tp = Py_TYPE(lz);
1876
14.0k
    PyObject_GC_UnTrack(lz);
1877
14.0k
    Py_XDECREF(lz->active);
1878
14.0k
    Py_XDECREF(lz->source);
1879
14.0k
    tp->tp_free(lz);
1880
14.0k
    Py_DECREF(tp);
1881
14.0k
}
1882
1883
static int
1884
chain_traverse(PyObject *op, visitproc visit, void *arg)
1885
13
{
1886
13
    chainobject *lz = chainobject_CAST(op);
1887
13
    Py_VISIT(Py_TYPE(lz));
1888
13
    Py_VISIT(lz->source);
1889
13
    Py_VISIT(lz->active);
1890
13
    return 0;
1891
13
}
1892
1893
static inline PyObject *
1894
chain_next_lock_held(PyObject *op)
1895
113k
{
1896
113k
    chainobject *lz = chainobject_CAST(op);
1897
113k
    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
141k
    while (lz->source != NULL) {
1903
141k
        if (lz->active == NULL) {
1904
42.0k
            PyObject *iterable = PyIter_Next(lz->source);
1905
42.0k
            if (iterable == NULL) {
1906
14.0k
                Py_CLEAR(lz->source);
1907
14.0k
                return NULL;            /* no more input sources */
1908
14.0k
            }
1909
28.0k
            lz->active = PyObject_GetIter(iterable);
1910
28.0k
            Py_DECREF(iterable);
1911
28.0k
            if (lz->active == NULL) {
1912
0
                Py_CLEAR(lz->source);
1913
0
                return NULL;            /* input not iterable */
1914
0
            }
1915
28.0k
        }
1916
127k
        item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1917
127k
        if (item != NULL)
1918
99.3k
            return item;
1919
28.0k
        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
28.0k
        Py_CLEAR(lz->active);
1927
28.0k
    }
1928
    /* Everything had been consumed already. */
1929
0
    return NULL;
1930
113k
}
1931
1932
static PyObject *
1933
chain_next(PyObject *op)
1934
113k
{
1935
113k
    PyObject *result;
1936
113k
    Py_BEGIN_CRITICAL_SECTION(op);
1937
113k
    result = chain_next_lock_held(op);
1938
113k
    Py_END_CRITICAL_SECTION()
1939
113k
    return result;
1940
113k
}
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
1.00k
#define productobject_CAST(op)  ((productobject *)(op))
1990
1991
static PyObject *
1992
product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1993
198
{
1994
198
    productobject *lz;
1995
198
    Py_ssize_t nargs, npools, repeat=1;
1996
198
    PyObject *pools = NULL;
1997
198
    Py_ssize_t *indices = NULL;
1998
198
    Py_ssize_t i;
1999
2000
198
    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
198
    assert(PyTuple_CheckExact(args));
2019
198
    if (repeat == 0) {
2020
0
        nargs = 0;
2021
198
    } else {
2022
198
        nargs = PyTuple_GET_SIZE(args);
2023
198
        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
198
    }
2028
198
    npools = nargs * repeat;
2029
2030
198
    indices = PyMem_New(Py_ssize_t, npools);
2031
198
    if (indices == NULL) {
2032
0
        PyErr_NoMemory();
2033
0
        goto error;
2034
0
    }
2035
2036
198
    pools = PyTuple_New(npools);
2037
198
    if (pools == NULL)
2038
0
        goto error;
2039
2040
504
    for (i=0; i < nargs ; ++i) {
2041
306
        PyObject *item = PyTuple_GET_ITEM(args, i);
2042
306
        PyObject *pool = PySequence_Tuple(item);
2043
306
        if (pool == NULL)
2044
0
            goto error;
2045
306
        PyTuple_SET_ITEM(pools, i, pool);
2046
306
        indices[i] = 0;
2047
306
    }
2048
198
    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
198
    lz = (productobject *)type->tp_alloc(type, 0);
2057
198
    if (lz == NULL)
2058
0
        goto error;
2059
2060
198
    lz->pools = pools;
2061
198
    lz->indices = indices;
2062
198
    lz->result = NULL;
2063
198
    lz->stopped = 0;
2064
2065
198
    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
198
}
2073
2074
static void
2075
product_dealloc(PyObject *op)
2076
198
{
2077
198
    productobject *lz = productobject_CAST(op);
2078
198
    PyTypeObject *tp = Py_TYPE(lz);
2079
198
    PyObject_GC_UnTrack(lz);
2080
198
    Py_XDECREF(lz->pools);
2081
198
    Py_XDECREF(lz->result);
2082
198
    PyMem_Free(lz->indices);
2083
198
    tp->tp_free(lz);
2084
198
    Py_DECREF(tp);
2085
198
}
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
810
{
2111
810
    productobject *lz = productobject_CAST(op);
2112
810
    PyObject *pool;
2113
810
    PyObject *elem;
2114
810
    PyObject *oldelem;
2115
810
    PyObject *pools = lz->pools;
2116
810
    PyObject *result = lz->result;
2117
810
    Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2118
810
    Py_ssize_t i;
2119
2120
810
    if (lz->stopped)
2121
0
        return NULL;
2122
2123
810
    if (result == NULL) {
2124
        /* On the first pass, return an initial tuple filled with the
2125
           first element from each pool. */
2126
198
        result = PyTuple_New(npools);
2127
198
        if (result == NULL)
2128
0
            goto empty;
2129
198
        lz->result = result;
2130
504
        for (i=0; i < npools; i++) {
2131
306
            pool = PyTuple_GET_ITEM(pools, i);
2132
306
            if (PyTuple_GET_SIZE(pool) == 0)
2133
0
                goto empty;
2134
306
            elem = PyTuple_GET_ITEM(pool, 0);
2135
306
            Py_INCREF(elem);
2136
306
            PyTuple_SET_ITEM(result, i, elem);
2137
306
        }
2138
612
    } else {
2139
612
        Py_ssize_t *indices = lz->indices;
2140
2141
        /* Copy the previous result tuple or re-use it if available */
2142
612
        if (!_PyObject_IsUniquelyReferenced(result)) {
2143
612
            PyObject *old_result = result;
2144
612
            result = PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
2145
612
            if (result == NULL)
2146
0
                goto empty;
2147
612
            lz->result = result;
2148
612
            Py_DECREF(old_result);
2149
612
        }
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
612
        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
1.02k
        for (i=npools-1 ; i >= 0 ; i--) {
2161
828
            pool = PyTuple_GET_ITEM(pools, i);
2162
828
            indices[i]++;
2163
828
            if (indices[i] == PyTuple_GET_SIZE(pool)) {
2164
                /* Roll-over and advance to next pool */
2165
414
                indices[i] = 0;
2166
414
                elem = PyTuple_GET_ITEM(pool, 0);
2167
414
                Py_INCREF(elem);
2168
414
                oldelem = PyTuple_GET_ITEM(result, i);
2169
414
                PyTuple_SET_ITEM(result, i, elem);
2170
414
                Py_DECREF(oldelem);
2171
414
            } else {
2172
                /* No rollover. Just increment and stop here. */
2173
414
                elem = PyTuple_GET_ITEM(pool, indices[i]);
2174
414
                Py_INCREF(elem);
2175
414
                oldelem = PyTuple_GET_ITEM(result, i);
2176
414
                PyTuple_SET_ITEM(result, i, elem);
2177
414
                Py_DECREF(oldelem);
2178
414
                break;
2179
414
            }
2180
828
        }
2181
2182
        /* If i is negative, then the indices have all rolled-over
2183
           and we're done. */
2184
612
        if (i < 0)
2185
198
            goto empty;
2186
612
    }
2187
2188
612
    return Py_NewRef(result);
2189
2190
198
empty:
2191
198
    lz->stopped = 1;
2192
198
    return NULL;
2193
810
}
2194
2195
static PyObject *
2196
product_next(PyObject *op)
2197
810
{
2198
810
    PyObject *result;
2199
810
    Py_BEGIN_CRITICAL_SECTION(op);
2200
810
    result = product_next_lock_held(op);
2201
810
    Py_END_CRITICAL_SECTION()
2202
810
    return result;
2203
810
}
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
486
#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
144
{
2771
144
    permutationsobject *po;
2772
144
    Py_ssize_t n;
2773
144
    Py_ssize_t r;
2774
144
    PyObject *pool = NULL;
2775
144
    Py_ssize_t *indices = NULL;
2776
144
    Py_ssize_t *cycles = NULL;
2777
144
    Py_ssize_t i;
2778
2779
144
    pool = PySequence_Tuple(iterable);
2780
144
    if (pool == NULL)
2781
0
        goto error;
2782
144
    n = PyTuple_GET_SIZE(pool);
2783
2784
144
    r = n;
2785
144
    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
144
    if (r < 0) {
2795
0
        PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2796
0
        goto error;
2797
0
    }
2798
2799
144
    indices = PyMem_New(Py_ssize_t, n);
2800
144
    cycles = PyMem_New(Py_ssize_t, r);
2801
144
    if (indices == NULL || cycles == NULL) {
2802
0
        PyErr_NoMemory();
2803
0
        goto error;
2804
0
    }
2805
2806
342
    for (i=0 ; i<n ; i++)
2807
198
        indices[i] = i;
2808
342
    for (i=0 ; i<r ; i++)
2809
198
        cycles[i] = n - i;
2810
2811
    /* create permutationsobject structure */
2812
144
    po = (permutationsobject *)type->tp_alloc(type, 0);
2813
144
    if (po == NULL)
2814
0
        goto error;
2815
2816
144
    po->pool = pool;
2817
144
    po->indices = indices;
2818
144
    po->cycles = cycles;
2819
144
    po->result = NULL;
2820
144
    po->r = r;
2821
144
    po->stopped = r > n ? 1 : 0;
2822
2823
144
    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
144
}
2833
2834
static void
2835
permutations_dealloc(PyObject *op)
2836
144
{
2837
144
    permutationsobject *po = permutationsobject_CAST(op);
2838
144
    PyTypeObject *tp = Py_TYPE(po);
2839
144
    PyObject_GC_UnTrack(po);
2840
144
    Py_XDECREF(po->pool);
2841
144
    Py_XDECREF(po->result);
2842
144
    PyMem_Free(po->indices);
2843
144
    PyMem_Free(po->cycles);
2844
144
    tp->tp_free(po);
2845
144
    Py_DECREF(tp);
2846
144
}
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
342
{
2871
342
    permutationsobject *po = permutationsobject_CAST(op);
2872
342
    PyObject *elem;
2873
342
    PyObject *oldelem;
2874
342
    PyObject *pool = po->pool;
2875
342
    Py_ssize_t *indices = po->indices;
2876
342
    Py_ssize_t *cycles = po->cycles;
2877
342
    PyObject *result = po->result;
2878
342
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2879
342
    Py_ssize_t r = po->r;
2880
342
    Py_ssize_t i, j, k, index;
2881
2882
342
    if (po->stopped)
2883
0
        return NULL;
2884
2885
342
    if (result == NULL) {
2886
        /* On the first pass, initialize result tuple using the indices */
2887
144
        result = PyTuple_New(r);
2888
144
        if (result == NULL)
2889
0
            goto empty;
2890
144
        po->result = result;
2891
342
        for (i=0; i<r ; i++) {
2892
198
            index = indices[i];
2893
198
            elem = PyTuple_GET_ITEM(pool, index);
2894
198
            Py_INCREF(elem);
2895
198
            PyTuple_SET_ITEM(result, i, elem);
2896
198
        }
2897
198
    } else {
2898
198
        if (n == 0)
2899
0
            goto empty;
2900
2901
        /* Copy the previous result tuple or re-use it if available */
2902
198
        if (!_PyObject_IsUniquelyReferenced(result)) {
2903
198
            PyObject *old_result = result;
2904
198
            result = PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2905
198
            if (result == NULL)
2906
0
                goto empty;
2907
198
            po->result = result;
2908
198
            Py_DECREF(old_result);
2909
198
        }
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
198
        assert(r == 0 || Py_REFCNT(result) == 1);
2917
2918
        /* Decrement rightmost cycle, moving leftward upon zero rollover */
2919
450
        for (i=r-1 ; i>=0 ; i--) {
2920
306
            cycles[i] -= 1;
2921
306
            if (cycles[i] == 0) {
2922
                /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2923
252
                index = indices[i];
2924
306
                for (j=i ; j<n-1 ; j++)
2925
54
                    indices[j] = indices[j+1];
2926
252
                indices[n-1] = index;
2927
252
                cycles[i] = n - i;
2928
252
            } else {
2929
54
                j = cycles[i];
2930
54
                index = indices[i];
2931
54
                indices[i] = indices[n-j];
2932
54
                indices[n-j] = index;
2933
2934
162
                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
108
                    index = indices[k];
2938
108
                    elem = PyTuple_GET_ITEM(pool, index);
2939
108
                    Py_INCREF(elem);
2940
108
                    oldelem = PyTuple_GET_ITEM(result, k);
2941
108
                    PyTuple_SET_ITEM(result, k, elem);
2942
108
                    Py_DECREF(oldelem);
2943
108
                }
2944
54
                break;
2945
54
            }
2946
306
        }
2947
        /* If i is negative, then the cycles have all
2948
           rolled-over and we're done. */
2949
198
        if (i < 0)
2950
144
            goto empty;
2951
198
    }
2952
198
    return Py_NewRef(result);
2953
2954
144
empty:
2955
144
    po->stopped = 1;
2956
144
    return NULL;
2957
342
}
2958
2959
static PyObject *
2960
permutations_next(PyObject *op)
2961
342
{
2962
342
    PyObject *result;
2963
342
    Py_BEGIN_CRITICAL_SECTION(op);
2964
342
    result = permutations_next_lock_held(op);
2965
342
    Py_END_CRITICAL_SECTION()
2966
342
    return result;
2967
342
}
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
282
#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
44
{
3290
44
    PyObject *it;
3291
44
    filterfalseobject *lz;
3292
3293
    /* Get iterator. */
3294
44
    it = PyObject_GetIter(seq);
3295
44
    if (it == NULL)
3296
0
        return NULL;
3297
3298
    /* create filterfalseobject structure */
3299
44
    lz = (filterfalseobject *)type->tp_alloc(type, 0);
3300
44
    if (lz == NULL) {
3301
0
        Py_DECREF(it);
3302
0
        return NULL;
3303
0
    }
3304
44
    lz->func = Py_NewRef(func);
3305
44
    lz->it = it;
3306
3307
44
    return (PyObject *)lz;
3308
44
}
3309
3310
static void
3311
filterfalse_dealloc(PyObject *op)
3312
44
{
3313
44
    filterfalseobject *lz = filterfalseobject_CAST(op);
3314
44
    PyTypeObject *tp = Py_TYPE(lz);
3315
44
    PyObject_GC_UnTrack(lz);
3316
44
    Py_XDECREF(lz->func);
3317
44
    Py_XDECREF(lz->it);
3318
44
    tp->tp_free(lz);
3319
44
    Py_DECREF(tp);
3320
44
}
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
238
{
3335
238
    filterfalseobject *lz = filterfalseobject_CAST(op);
3336
238
    PyObject *item;
3337
238
    PyObject *it = lz->it;
3338
238
    long ok;
3339
238
    PyObject *(*iternext)(PyObject *);
3340
3341
238
    iternext = *Py_TYPE(it)->tp_iternext;
3342
250
    for (;;) {
3343
250
        item = iternext(it);
3344
250
        if (item == NULL)
3345
44
            return NULL;
3346
3347
206
        if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3348
0
            ok = PyObject_IsTrue(item);
3349
206
        } else {
3350
206
            PyObject *good;
3351
206
            good = PyObject_CallOneArg(lz->func, item);
3352
206
            if (good == NULL) {
3353
0
                Py_DECREF(item);
3354
0
                return NULL;
3355
0
            }
3356
206
            ok = PyObject_IsTrue(good);
3357
206
            Py_DECREF(good);
3358
206
        }
3359
206
        if (ok == 0)
3360
194
            return item;
3361
12
        Py_DECREF(item);
3362
12
        if (ok < 0)
3363
0
            return NULL;
3364
12
    }
3365
238
}
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
5.05k
#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
20
{
3436
20
    countobject *lz;
3437
20
    int fast_mode;
3438
20
    Py_ssize_t cnt = 0;
3439
20
    long step;
3440
3441
20
    if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3442
20
        (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
20
    fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3448
20
                (long_step == NULL || PyLong_Check(long_step));
3449
3450
    /* If not specified, start defaults to 0 */
3451
20
    if (long_cnt != NULL) {
3452
6
        if (fast_mode) {
3453
6
            assert(PyLong_Check(long_cnt));
3454
6
            cnt = PyLong_AsSsize_t(long_cnt);
3455
6
            if (cnt == -1 && PyErr_Occurred()) {
3456
0
                PyErr_Clear();
3457
0
                fast_mode = 0;
3458
0
            }
3459
6
        }
3460
14
    } else {
3461
14
        cnt = 0;
3462
14
        long_cnt = _PyLong_GetZero();
3463
14
    }
3464
20
    Py_INCREF(long_cnt);
3465
3466
    /* If not specified, step defaults to 1 */
3467
20
    if (long_step == NULL) {
3468
20
        long_step = _PyLong_GetOne();
3469
20
    }
3470
20
    Py_INCREF(long_step);
3471
3472
20
    assert(long_cnt != NULL && long_step != NULL);
3473
3474
    /* Fast mode only works when the step is 1 */
3475
20
    if (fast_mode) {
3476
20
        assert(PyLong_Check(long_step));
3477
20
        step = PyLong_AsLong(long_step);
3478
20
        if (step != 1) {
3479
0
            fast_mode = 0;
3480
0
            if (step == -1 && PyErr_Occurred())
3481
0
                PyErr_Clear();
3482
0
        }
3483
20
    }
3484
3485
20
    if (fast_mode)
3486
20
        Py_CLEAR(long_cnt);
3487
0
    else
3488
0
        cnt = PY_SSIZE_T_MAX;
3489
3490
20
    assert((long_cnt == NULL && fast_mode) ||
3491
20
           (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
3492
20
    assert(!fast_mode ||
3493
20
           (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
3494
3495
    /* create countobject structure */
3496
20
    lz = (countobject *)type->tp_alloc(type, 0);
3497
20
    if (lz == NULL) {
3498
0
        Py_XDECREF(long_cnt);
3499
0
        Py_DECREF(long_step);
3500
0
        return NULL;
3501
0
    }
3502
20
    lz->cnt = cnt;
3503
20
    lz->long_cnt = long_cnt;
3504
20
    lz->long_step = long_step;
3505
3506
20
    return (PyObject *)lz;
3507
20
}
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
317
{
3524
317
    countobject *lz = countobject_CAST(op);
3525
317
    Py_VISIT(Py_TYPE(lz));
3526
317
    Py_VISIT(lz->long_cnt);
3527
317
    Py_VISIT(lz->long_step);
3528
317
    return 0;
3529
317
}
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
4.73k
{
3556
4.73k
    countobject *lz = countobject_CAST(op);
3557
4.73k
#ifndef Py_GIL_DISABLED
3558
4.73k
    if (lz->cnt == PY_SSIZE_T_MAX)
3559
0
        return count_nextlong(lz);
3560
4.73k
    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
4.73k
}
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
47.4k
#define repeatobject_CAST(op)   ((repeatobject *)(op))
3639
3640
static PyObject *
3641
repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3642
4.74k
{
3643
4.74k
    repeatobject *ro;
3644
4.74k
    PyObject *element;
3645
4.74k
    Py_ssize_t cnt = -1, n_args;
3646
4.74k
    static char *kwargs[] = {"object", "times", NULL};
3647
3648
4.74k
    n_args = PyTuple_GET_SIZE(args);
3649
4.74k
    if (kwds != NULL)
3650
0
        n_args += PyDict_GET_SIZE(kwds);
3651
4.74k
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3652
4.74k
                                     &element, &cnt))
3653
0
        return NULL;
3654
    /* Does user supply times argument? */
3655
4.74k
    if (n_args == 2 && cnt < 0)
3656
0
        cnt = 0;
3657
3658
4.74k
    ro = (repeatobject *)type->tp_alloc(type, 0);
3659
4.74k
    if (ro == NULL)
3660
0
        return NULL;
3661
4.74k
    ro->element = Py_NewRef(element);
3662
4.74k
    ro->cnt = cnt;
3663
4.74k
    return (PyObject *)ro;
3664
4.74k
}
3665
3666
static void
3667
repeat_dealloc(PyObject *op)
3668
4.74k
{
3669
4.74k
    repeatobject *ro = repeatobject_CAST(op);
3670
4.74k
    PyTypeObject *tp = Py_TYPE(ro);
3671
4.74k
    PyObject_GC_UnTrack(ro);
3672
4.74k
    Py_XDECREF(ro->element);
3673
4.74k
    tp->tp_free(ro);
3674
4.74k
    Py_DECREF(tp);
3675
4.74k
}
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
42.6k
{
3689
42.6k
    repeatobject *ro = repeatobject_CAST(op);
3690
42.6k
    Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(ro->cnt);
3691
42.6k
    if (cnt == 0) {
3692
4.74k
        return NULL;
3693
4.74k
    }
3694
37.9k
    if (cnt > 0) {
3695
37.9k
        cnt--;
3696
37.9k
        FT_ATOMIC_STORE_SSIZE_RELAXED(ro->cnt, cnt);
3697
37.9k
    }
3698
37.9k
    return Py_NewRef(ro->element);
3699
42.6k
}
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
1.51k
{
4007
1.51k
    itertools_state *state = get_module_state(mod);
4008
1.51k
    Py_VISIT(state->accumulate_type);
4009
1.51k
    Py_VISIT(state->batched_type);
4010
1.51k
    Py_VISIT(state->chain_type);
4011
1.51k
    Py_VISIT(state->combinations_type);
4012
1.51k
    Py_VISIT(state->compress_type);
4013
1.51k
    Py_VISIT(state->count_type);
4014
1.51k
    Py_VISIT(state->cwr_type);
4015
1.51k
    Py_VISIT(state->cycle_type);
4016
1.51k
    Py_VISIT(state->dropwhile_type);
4017
1.51k
    Py_VISIT(state->filterfalse_type);
4018
1.51k
    Py_VISIT(state->groupby_type);
4019
1.51k
    Py_VISIT(state->_grouper_type);
4020
1.51k
    Py_VISIT(state->islice_type);
4021
1.51k
    Py_VISIT(state->pairwise_type);
4022
1.51k
    Py_VISIT(state->permutations_type);
4023
1.51k
    Py_VISIT(state->product_type);
4024
1.51k
    Py_VISIT(state->repeat_type);
4025
1.51k
    Py_VISIT(state->starmap_type);
4026
1.51k
    Py_VISIT(state->takewhile_type);
4027
1.51k
    Py_VISIT(state->tee_type);
4028
1.51k
    Py_VISIT(state->teedataobject_type);
4029
1.51k
    Py_VISIT(state->ziplongest_type);
4030
1.51k
    return 0;
4031
1.51k
}
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
616
#define ADD_TYPE(module, type, spec)                                     \
4069
616
do {                                                                     \
4070
616
    type = (PyTypeObject *)PyType_FromModuleAndSpec(module, spec, NULL); \
4071
616
    if (type == NULL) {                                                  \
4072
0
        return -1;                                                       \
4073
0
    }                                                                    \
4074
616
    if (PyModule_AddType(module, type) < 0) {                            \
4075
0
        return -1;                                                       \
4076
0
    }                                                                    \
4077
616
} while (0)
4078
4079
static int
4080
itertoolsmodule_exec(PyObject *mod)
4081
28
{
4082
28
    itertools_state *state = get_module_state(mod);
4083
28
    ADD_TYPE(mod, state->accumulate_type, &accumulate_spec);
4084
28
    ADD_TYPE(mod, state->batched_type, &batched_spec);
4085
28
    ADD_TYPE(mod, state->chain_type, &chain_spec);
4086
28
    ADD_TYPE(mod, state->combinations_type, &combinations_spec);
4087
28
    ADD_TYPE(mod, state->compress_type, &compress_spec);
4088
28
    ADD_TYPE(mod, state->count_type, &count_spec);
4089
28
    ADD_TYPE(mod, state->cwr_type, &cwr_spec);
4090
28
    ADD_TYPE(mod, state->cycle_type, &cycle_spec);
4091
28
    ADD_TYPE(mod, state->dropwhile_type, &dropwhile_spec);
4092
28
    ADD_TYPE(mod, state->filterfalse_type, &filterfalse_spec);
4093
28
    ADD_TYPE(mod, state->groupby_type, &groupby_spec);
4094
28
    ADD_TYPE(mod, state->_grouper_type, &_grouper_spec);
4095
28
    ADD_TYPE(mod, state->islice_type, &islice_spec);
4096
28
    ADD_TYPE(mod, state->pairwise_type, &pairwise_spec);
4097
28
    ADD_TYPE(mod, state->permutations_type, &permutations_spec);
4098
28
    ADD_TYPE(mod, state->product_type, &product_spec);
4099
28
    ADD_TYPE(mod, state->repeat_type, &repeat_spec);
4100
28
    ADD_TYPE(mod, state->starmap_type, &starmap_spec);
4101
28
    ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
4102
28
    ADD_TYPE(mod, state->tee_type, &tee_spec);
4103
28
    ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
4104
28
    ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
4105
4106
28
    Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
4107
28
    return 0;
4108
28
}
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
28
{
4138
28
    return PyModuleDef_Init(&itertoolsmodule);
4139
28
}