Coverage Report

Created: 2026-05-30 06:18

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