Coverage Report

Created: 2026-04-12 06:54

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