Coverage Report

Created: 2026-06-07 06:26

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/src/cpython3/Modules/itertoolsmodule.c
Line
Count
Source
1
#include "Python.h"
2
#include "pycore_call.h"              // _PyObject_CallNoArgs()
3
#include "pycore_ceval.h"             // _PyEval_GetBuiltin()
4
#include "pycore_critical_section.h"  // Py_BEGIN_CRITICAL_SECTION()
5
#include "pycore_long.h"              // _PyLong_GetZero()
6
#include "pycore_moduleobject.h"      // _PyModule_GetState()
7
#include "pycore_typeobject.h"        // _PyType_GetModuleState()
8
#include "pycore_object.h"            // _PyObject_GC_TRACK()
9
#include "pycore_tuple.h"             // _PyTuple_ITEMS()
10
11
#include <stddef.h>                   // offsetof()
12
13
/* Itertools module written and maintained
14
   by Raymond D. Hettinger <python@rcn.com>
15
*/
16
17
typedef struct {
18
    PyTypeObject *accumulate_type;
19
    PyTypeObject *batched_type;
20
    PyTypeObject *chain_type;
21
    PyTypeObject *combinations_type;
22
    PyTypeObject *compress_type;
23
    PyTypeObject *count_type;
24
    PyTypeObject *cwr_type;
25
    PyTypeObject *cycle_type;
26
    PyTypeObject *dropwhile_type;
27
    PyTypeObject *filterfalse_type;
28
    PyTypeObject *groupby_type;
29
    PyTypeObject *_grouper_type;
30
    PyTypeObject *islice_type;
31
    PyTypeObject *pairwise_type;
32
    PyTypeObject *permutations_type;
33
    PyTypeObject *product_type;
34
    PyTypeObject *repeat_type;
35
    PyTypeObject *starmap_type;
36
    PyTypeObject *takewhile_type;
37
    PyTypeObject *tee_type;
38
    PyTypeObject *teedataobject_type;
39
    PyTypeObject *ziplongest_type;
40
} itertools_state;
41
42
static inline itertools_state *
43
get_module_state(PyObject *mod)
44
65
{
45
65
    void *state = _PyModule_GetState(mod);
46
65
    assert(state != NULL);
47
65
    return (itertools_state *)state;
48
65
}
49
50
static inline itertools_state *
51
get_module_state_by_cls(PyTypeObject *cls)
52
0
{
53
0
    void *state = _PyType_GetModuleState(cls);
54
0
    assert(state != NULL);
55
0
    return (itertools_state *)state;
56
0
}
57
58
static struct PyModuleDef itertoolsmodule;
59
60
static inline itertools_state *
61
find_state_by_type(PyTypeObject *tp)
62
0
{
63
0
    PyObject *mod = PyType_GetModuleByDef(tp, &itertoolsmodule);
64
0
    assert(mod != NULL);
65
0
    return get_module_state(mod);
66
0
}
67
68
/*[clinic input]
69
module itertools
70
class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
71
class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
72
class itertools.teedataobject "teedataobject *" "clinic_state()->teedataobject_type"
73
class itertools._tee "teeobject *" "clinic_state()->tee_type"
74
class itertools.batched "batchedobject *" "clinic_state()->batched_type"
75
class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
76
class itertools.dropwhile "dropwhileobject *" "clinic_state()->dropwhile_type"
77
class itertools.takewhile "takewhileobject *" "clinic_state()->takewhile_type"
78
class itertools.starmap "starmapobject *" "clinic_state()->starmap_type"
79
class itertools.chain "chainobject *" "clinic_state()->chain_type"
80
class itertools.combinations "combinationsobject *" "clinic_state()->combinations_type"
81
class itertools.combinations_with_replacement "cwr_object *" "clinic_state()->cwr_type"
82
class itertools.permutations "permutationsobject *" "clinic_state()->permutations_type"
83
class itertools.accumulate "accumulateobject *" "clinic_state()->accumulate_type"
84
class itertools.compress "compressobject *" "clinic_state()->compress_type"
85
class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_type"
86
class itertools.count "countobject *" "clinic_state()->count_type"
87
class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
88
[clinic start generated code]*/
89
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=aa48fe4de9d4080f]*/
90
91
0
#define clinic_state() (find_state_by_type(type))
92
0
#define clinic_state_by_cls() (get_module_state_by_cls(base_tp))
93
#include "clinic/itertoolsmodule.c.h"
94
#undef clinic_state_by_cls
95
#undef clinic_state
96
97
98
/* batched object ************************************************************/
99
100
typedef struct {
101
    PyObject_HEAD
102
    PyObject *it;
103
    Py_ssize_t batch_size;
104
    bool strict;
105
} batchedobject;
106
107
0
#define batchedobject_CAST(op)  ((batchedobject *)(op))
108
109
/*[clinic input]
110
@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
0
#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
0
{
300
0
    PyObject *it;
301
0
    pairwiseobject *po;
302
303
0
    it = PyObject_GetIter(iterable);
304
0
    if (it == NULL) {
305
0
        return NULL;
306
0
    }
307
0
    po = (pairwiseobject *)type->tp_alloc(type, 0);
308
0
    if (po == NULL) {
309
0
        Py_DECREF(it);
310
0
        return NULL;
311
0
    }
312
0
    po->it = it;
313
0
    po->old = NULL;
314
0
    po->result = _PyTuple_FromPairSteal(Py_None, Py_None);
315
0
    if (po->result == NULL) {
316
0
        Py_DECREF(po);
317
0
        return NULL;
318
0
    }
319
0
    return (PyObject *)po;
320
0
}
321
322
static void
323
pairwise_dealloc(PyObject *op)
324
0
{
325
0
    pairwiseobject *po = pairwiseobject_CAST(op);
326
0
    PyTypeObject *tp = Py_TYPE(po);
327
0
    PyObject_GC_UnTrack(po);
328
0
    Py_XDECREF(po->it);
329
0
    Py_XDECREF(po->old);
330
0
    Py_XDECREF(po->result);
331
0
    tp->tp_free(po);
332
0
    Py_DECREF(tp);
333
0
}
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
0
{
349
0
    pairwiseobject *po = pairwiseobject_CAST(op);
350
0
    PyObject *it = po->it;
351
0
    PyObject *old = po->old;
352
0
    PyObject *new, *result;
353
354
0
    if (it == NULL) {
355
0
        return NULL;
356
0
    }
357
0
    if (old == NULL) {
358
0
        old = (*Py_TYPE(it)->tp_iternext)(it);
359
0
        Py_XSETREF(po->old, old);
360
0
        if (old == NULL) {
361
0
            Py_CLEAR(po->it);
362
0
            return NULL;
363
0
        }
364
0
        it = po->it;
365
0
        if (it == NULL) {
366
0
            Py_CLEAR(po->old);
367
0
            return NULL;
368
0
        }
369
0
    }
370
0
    Py_INCREF(old);
371
0
    new = (*Py_TYPE(it)->tp_iternext)(it);
372
0
    if (new == NULL) {
373
0
        Py_CLEAR(po->it);
374
0
        Py_CLEAR(po->old);
375
0
        Py_DECREF(old);
376
0
        return NULL;
377
0
    }
378
379
0
    result = po->result;
380
0
    if (_PyObject_IsUniquelyReferenced(result)) {
381
0
        Py_INCREF(result);
382
0
        PyObject *last_old = PyTuple_GET_ITEM(result, 0);
383
0
        PyObject *last_new = PyTuple_GET_ITEM(result, 1);
384
0
        PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
385
0
        PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
386
0
        Py_DECREF(last_old);
387
0
        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
0
        _PyTuple_Recycle(result);
391
0
    }
392
0
    else {
393
0
        result = _PyTuple_FromPair(old, new);
394
0
    }
395
396
0
    Py_XSETREF(po->old, new);
397
0
    Py_DECREF(old);
398
0
    return result;
399
0
}
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
0
#define chainobject_CAST(op)    ((chainobject *)(op))
1819
1820
static PyObject *
1821
chain_new_internal(PyTypeObject *type, PyObject *source)
1822
0
{
1823
0
    chainobject *lz;
1824
1825
0
    lz = (chainobject *)type->tp_alloc(type, 0);
1826
0
    if (lz == NULL) {
1827
0
        Py_DECREF(source);
1828
0
        return NULL;
1829
0
    }
1830
1831
0
    lz->source = source;
1832
0
    lz->active = NULL;
1833
0
    return (PyObject *)lz;
1834
0
}
1835
1836
static PyObject *
1837
chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1838
0
{
1839
0
    PyObject *source;
1840
1841
0
    itertools_state *state = find_state_by_type(type);
1842
0
    PyTypeObject *chain_type = state->chain_type;
1843
0
    if ((type == chain_type || type->tp_init == chain_type->tp_init) &&
1844
0
        !_PyArg_NoKeywords("chain", kwds))
1845
0
        return NULL;
1846
1847
0
    source = PyObject_GetIter(args);
1848
0
    if (source == NULL)
1849
0
        return NULL;
1850
1851
0
    return chain_new_internal(type, source);
1852
0
}
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
0
{
1867
0
    PyObject *source;
1868
1869
0
    source = PyObject_GetIter(arg);
1870
0
    if (source == NULL)
1871
0
        return NULL;
1872
1873
0
    return chain_new_internal(type, source);
1874
0
}
1875
1876
static void
1877
chain_dealloc(PyObject *op)
1878
0
{
1879
0
    chainobject *lz = chainobject_CAST(op);
1880
0
    PyTypeObject *tp = Py_TYPE(lz);
1881
0
    PyObject_GC_UnTrack(lz);
1882
0
    Py_XDECREF(lz->active);
1883
0
    Py_XDECREF(lz->source);
1884
0
    tp->tp_free(lz);
1885
0
    Py_DECREF(tp);
1886
0
}
1887
1888
static int
1889
chain_traverse(PyObject *op, visitproc visit, void *arg)
1890
0
{
1891
0
    chainobject *lz = chainobject_CAST(op);
1892
0
    Py_VISIT(Py_TYPE(lz));
1893
0
    Py_VISIT(lz->source);
1894
0
    Py_VISIT(lz->active);
1895
0
    return 0;
1896
0
}
1897
1898
static inline PyObject *
1899
chain_next_lock_held(PyObject *op)
1900
0
{
1901
0
    chainobject *lz = chainobject_CAST(op);
1902
0
    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
0
    while (lz->source != NULL) {
1908
0
        if (lz->active == NULL) {
1909
0
            PyObject *iterable = PyIter_Next(lz->source);
1910
0
            if (iterable == NULL) {
1911
0
                Py_CLEAR(lz->source);
1912
0
                return NULL;            /* no more input sources */
1913
0
            }
1914
0
            lz->active = PyObject_GetIter(iterable);
1915
0
            Py_DECREF(iterable);
1916
0
            if (lz->active == NULL) {
1917
0
                Py_CLEAR(lz->source);
1918
0
                return NULL;            /* input not iterable */
1919
0
            }
1920
0
        }
1921
0
        item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1922
0
        if (item != NULL)
1923
0
            return item;
1924
0
        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
0
        Py_CLEAR(lz->active);
1932
0
    }
1933
    /* Everything had been consumed already. */
1934
0
    return NULL;
1935
0
}
1936
1937
static PyObject *
1938
chain_next(PyObject *op)
1939
0
{
1940
0
    PyObject *result;
1941
0
    Py_BEGIN_CRITICAL_SECTION(op);
1942
0
    result = chain_next_lock_held(op);
1943
0
    Py_END_CRITICAL_SECTION()
1944
0
    return result;
1945
0
}
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
PyDoc_STRVAR(chain_class_getitem_doc,
1956
"chain is generic over the type of its contents.\n\
1957
This is the union of the types of the input iterable contents.");
1958
1959
static PyMethodDef chain_methods[] = {
1960
    ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
1961
    {"__class_getitem__",    Py_GenericAlias,
1962
    METH_O|METH_CLASS,       chain_class_getitem_doc},
1963
    {NULL,              NULL}           /* sentinel */
1964
};
1965
1966
static PyType_Slot chain_slots[] = {
1967
    {Py_tp_dealloc, chain_dealloc},
1968
    {Py_tp_getattro, PyObject_GenericGetAttr},
1969
    {Py_tp_doc, (void *)chain_doc},
1970
    {Py_tp_traverse, chain_traverse},
1971
    {Py_tp_iter, PyObject_SelfIter},
1972
    {Py_tp_iternext, chain_next},
1973
    {Py_tp_methods, chain_methods},
1974
    {Py_tp_new, chain_new},
1975
    {Py_tp_free, PyObject_GC_Del},
1976
    {0, NULL},
1977
};
1978
1979
static PyType_Spec chain_spec = {
1980
    .name = "itertools.chain",
1981
    .basicsize = sizeof(chainobject),
1982
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1983
              Py_TPFLAGS_IMMUTABLETYPE),
1984
    .slots = chain_slots,
1985
};
1986
1987
1988
/* product object ************************************************************/
1989
1990
typedef struct {
1991
    PyObject_HEAD
1992
    PyObject *pools;        /* tuple of pool tuples */
1993
    Py_ssize_t *indices;    /* one index per pool */
1994
    PyObject *result;       /* most recently returned result tuple */
1995
    int stopped;            /* set to 1 when the iterator is exhausted */
1996
} productobject;
1997
1998
336
#define productobject_CAST(op)  ((productobject *)(op))
1999
2000
static PyObject *
2001
product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2002
66
{
2003
66
    productobject *lz;
2004
66
    Py_ssize_t nargs, npools, repeat=1;
2005
66
    PyObject *pools = NULL;
2006
66
    Py_ssize_t *indices = NULL;
2007
66
    Py_ssize_t i;
2008
2009
66
    if (kwds != NULL) {
2010
0
        char *kwlist[] = {"repeat", 0};
2011
0
        PyObject *tmpargs = PyTuple_New(0);
2012
0
        if (tmpargs == NULL)
2013
0
            return NULL;
2014
0
        if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2015
0
                                         kwlist, &repeat)) {
2016
0
            Py_DECREF(tmpargs);
2017
0
            return NULL;
2018
0
        }
2019
0
        Py_DECREF(tmpargs);
2020
0
        if (repeat < 0) {
2021
0
            PyErr_SetString(PyExc_ValueError,
2022
0
                            "repeat argument cannot be negative");
2023
0
            return NULL;
2024
0
        }
2025
0
    }
2026
2027
66
    assert(PyTuple_CheckExact(args));
2028
66
    if (repeat == 0) {
2029
0
        nargs = 0;
2030
66
    } else {
2031
66
        nargs = PyTuple_GET_SIZE(args);
2032
66
        if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
2033
0
            PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2034
0
            return NULL;
2035
0
        }
2036
66
    }
2037
66
    npools = nargs * repeat;
2038
2039
66
    indices = PyMem_New(Py_ssize_t, npools);
2040
66
    if (indices == NULL) {
2041
0
        PyErr_NoMemory();
2042
0
        goto error;
2043
0
    }
2044
2045
66
    pools = PyTuple_New(npools);
2046
66
    if (pools == NULL)
2047
0
        goto error;
2048
2049
168
    for (i=0; i < nargs ; ++i) {
2050
102
        PyObject *item = PyTuple_GET_ITEM(args, i);
2051
0
        PyObject *pool = PySequence_Tuple(item);
2052
102
        if (pool == NULL)
2053
0
            goto error;
2054
102
        PyTuple_SET_ITEM(pools, i, pool);
2055
102
        indices[i] = 0;
2056
102
    }
2057
66
    for ( ; i < npools; ++i) {
2058
0
        PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2059
0
        Py_INCREF(pool);
2060
0
        PyTuple_SET_ITEM(pools, i, pool);
2061
0
        indices[i] = 0;
2062
0
    }
2063
2064
    /* create productobject structure */
2065
66
    lz = (productobject *)type->tp_alloc(type, 0);
2066
66
    if (lz == NULL)
2067
0
        goto error;
2068
2069
66
    lz->pools = pools;
2070
66
    lz->indices = indices;
2071
66
    lz->result = NULL;
2072
66
    lz->stopped = 0;
2073
2074
66
    return (PyObject *)lz;
2075
2076
0
error:
2077
0
    if (indices != NULL)
2078
0
        PyMem_Free(indices);
2079
0
    Py_XDECREF(pools);
2080
0
    return NULL;
2081
66
}
2082
2083
static void
2084
product_dealloc(PyObject *op)
2085
66
{
2086
66
    productobject *lz = productobject_CAST(op);
2087
66
    PyTypeObject *tp = Py_TYPE(lz);
2088
66
    PyObject_GC_UnTrack(lz);
2089
66
    Py_XDECREF(lz->pools);
2090
66
    Py_XDECREF(lz->result);
2091
66
    PyMem_Free(lz->indices);
2092
66
    tp->tp_free(lz);
2093
66
    Py_DECREF(tp);
2094
66
}
2095
2096
static PyObject *
2097
product_sizeof(PyObject *op, PyObject *Py_UNUSED(ignored))
2098
0
{
2099
0
    productobject *lz = productobject_CAST(op);
2100
0
    size_t res = _PyObject_SIZE(Py_TYPE(lz));
2101
0
    res += (size_t)PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2102
0
    return PyLong_FromSize_t(res);
2103
0
}
2104
2105
PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2106
2107
static int
2108
product_traverse(PyObject *op, visitproc visit, void *arg)
2109
0
{
2110
0
    productobject *lz = productobject_CAST(op);
2111
0
    Py_VISIT(Py_TYPE(lz));
2112
0
    Py_VISIT(lz->pools);
2113
0
    Py_VISIT(lz->result);
2114
0
    return 0;
2115
0
}
2116
2117
static PyObject *
2118
product_next_lock_held(PyObject *op)
2119
270
{
2120
270
    productobject *lz = productobject_CAST(op);
2121
270
    PyObject *pool;
2122
270
    PyObject *elem;
2123
270
    PyObject *oldelem;
2124
270
    PyObject *pools = lz->pools;
2125
270
    PyObject *result = lz->result;
2126
270
    Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2127
270
    Py_ssize_t i;
2128
2129
270
    if (lz->stopped)
2130
0
        return NULL;
2131
2132
270
    if (result == NULL) {
2133
        /* On the first pass, return an initial tuple filled with the
2134
           first element from each pool. */
2135
66
        result = PyTuple_New(npools);
2136
66
        if (result == NULL)
2137
0
            goto empty;
2138
66
        lz->result = result;
2139
168
        for (i=0; i < npools; i++) {
2140
102
            pool = PyTuple_GET_ITEM(pools, i);
2141
102
            if (PyTuple_GET_SIZE(pool) == 0)
2142
0
                goto empty;
2143
102
            elem = PyTuple_GET_ITEM(pool, 0);
2144
102
            Py_INCREF(elem);
2145
102
            PyTuple_SET_ITEM(result, i, elem);
2146
102
        }
2147
204
    } else {
2148
204
        Py_ssize_t *indices = lz->indices;
2149
2150
        /* Copy the previous result tuple or re-use it if available */
2151
204
        if (!_PyObject_IsUniquelyReferenced(result)) {
2152
204
            PyObject *old_result = result;
2153
204
            result = PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
2154
204
            if (result == NULL)
2155
0
                goto empty;
2156
204
            lz->result = result;
2157
204
            Py_DECREF(old_result);
2158
204
        }
2159
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2160
        // recycling it, make sure it's tracked again:
2161
0
        else {
2162
0
            _PyTuple_Recycle(result);
2163
0
        }
2164
        /* Now, we've got the only copy so we can update it in-place */
2165
204
        assert (npools==0 || Py_REFCNT(result) == 1);
2166
2167
        /* Update the pool indices right-to-left.  Only advance to the
2168
           next pool when the previous one rolls-over */
2169
342
        for (i=npools-1 ; i >= 0 ; i--) {
2170
276
            pool = PyTuple_GET_ITEM(pools, i);
2171
0
            indices[i]++;
2172
276
            if (indices[i] == PyTuple_GET_SIZE(pool)) {
2173
                /* Roll-over and advance to next pool */
2174
138
                indices[i] = 0;
2175
138
                elem = PyTuple_GET_ITEM(pool, 0);
2176
138
                Py_INCREF(elem);
2177
138
                oldelem = PyTuple_GET_ITEM(result, i);
2178
138
                PyTuple_SET_ITEM(result, i, elem);
2179
138
                Py_DECREF(oldelem);
2180
138
            } else {
2181
                /* No rollover. Just increment and stop here. */
2182
138
                elem = PyTuple_GET_ITEM(pool, indices[i]);
2183
138
                Py_INCREF(elem);
2184
138
                oldelem = PyTuple_GET_ITEM(result, i);
2185
138
                PyTuple_SET_ITEM(result, i, elem);
2186
138
                Py_DECREF(oldelem);
2187
138
                break;
2188
138
            }
2189
276
        }
2190
2191
        /* If i is negative, then the indices have all rolled-over
2192
           and we're done. */
2193
204
        if (i < 0)
2194
66
            goto empty;
2195
204
    }
2196
2197
204
    return Py_NewRef(result);
2198
2199
66
empty:
2200
66
    lz->stopped = 1;
2201
66
    return NULL;
2202
270
}
2203
2204
static PyObject *
2205
product_next(PyObject *op)
2206
270
{
2207
270
    PyObject *result;
2208
270
    Py_BEGIN_CRITICAL_SECTION(op);
2209
270
    result = product_next_lock_held(op);
2210
270
    Py_END_CRITICAL_SECTION()
2211
270
    return result;
2212
270
}
2213
2214
static PyMethodDef product_methods[] = {
2215
    {"__sizeof__", product_sizeof, METH_NOARGS, sizeof_doc},
2216
    {NULL,              NULL}   /* sentinel */
2217
};
2218
2219
PyDoc_STRVAR(product_doc,
2220
"product(*iterables, repeat=1)\n\
2221
--\n\
2222
\n\
2223
Cartesian product of input iterables.  Equivalent to nested for-loops.\n\n\
2224
For example, product(A, B) returns the same as:  ((x,y) for x in A for y in B).\n\
2225
The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2226
cycle in a manner similar to an odometer (with the rightmost element changing\n\
2227
on every iteration).\n\n\
2228
To compute the product of an iterable with itself, specify the number\n\
2229
of repetitions with the optional repeat keyword argument. For example,\n\
2230
product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
2231
product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2232
product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2233
2234
static PyType_Slot product_slots[] = {
2235
    {Py_tp_dealloc, product_dealloc},
2236
    {Py_tp_getattro, PyObject_GenericGetAttr},
2237
    {Py_tp_doc, (void *)product_doc},
2238
    {Py_tp_traverse, product_traverse},
2239
    {Py_tp_iter, PyObject_SelfIter},
2240
    {Py_tp_iternext, product_next},
2241
    {Py_tp_methods, product_methods},
2242
    {Py_tp_new, product_new},
2243
    {Py_tp_free, PyObject_GC_Del},
2244
    {0, NULL},
2245
};
2246
2247
static PyType_Spec product_spec = {
2248
    .name = "itertools.product",
2249
    .basicsize = sizeof(productobject),
2250
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2251
              Py_TPFLAGS_IMMUTABLETYPE),
2252
    .slots = product_slots,
2253
};
2254
2255
2256
/* combinations object *******************************************************/
2257
2258
typedef struct {
2259
    PyObject_HEAD
2260
    PyObject *pool;         /* input converted to a tuple */
2261
    Py_ssize_t *indices;    /* one index per result element */
2262
    PyObject *result;       /* most recently returned result tuple */
2263
    Py_ssize_t r;           /* size of result tuple */
2264
    int stopped;            /* set to 1 when the iterator is exhausted */
2265
} combinationsobject;
2266
2267
0
#define combinationsobject_CAST(op) ((combinationsobject *)(op))
2268
2269
/*[clinic input]
2270
@classmethod
2271
itertools.combinations.__new__
2272
    iterable: object
2273
    r: Py_ssize_t(allow_negative=False)
2274
Return successive r-length combinations of elements in the iterable.
2275
2276
combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2277
[clinic start generated code]*/
2278
2279
static PyObject *
2280
itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2281
                            Py_ssize_t r)
2282
/*[clinic end generated code: output=87a689b39c40039c input=a32f07a15cfa4676]*/
2283
0
{
2284
0
    combinationsobject *co;
2285
0
    Py_ssize_t n;
2286
0
    PyObject *pool = NULL;
2287
0
    Py_ssize_t *indices = NULL;
2288
0
    Py_ssize_t i;
2289
2290
0
    pool = PySequence_Tuple(iterable);
2291
0
    if (pool == NULL)
2292
0
        goto error;
2293
0
    n = PyTuple_GET_SIZE(pool);
2294
2295
0
    indices = PyMem_New(Py_ssize_t, r);
2296
0
    if (indices == NULL) {
2297
0
        PyErr_NoMemory();
2298
0
        goto error;
2299
0
    }
2300
2301
0
    for (i=0 ; i<r ; i++)
2302
0
        indices[i] = i;
2303
2304
    /* create combinationsobject structure */
2305
0
    co = (combinationsobject *)type->tp_alloc(type, 0);
2306
0
    if (co == NULL)
2307
0
        goto error;
2308
2309
0
    co->pool = pool;
2310
0
    co->indices = indices;
2311
0
    co->result = NULL;
2312
0
    co->r = r;
2313
0
    co->stopped = r > n ? 1 : 0;
2314
2315
0
    return (PyObject *)co;
2316
2317
0
error:
2318
0
    if (indices != NULL)
2319
0
        PyMem_Free(indices);
2320
0
    Py_XDECREF(pool);
2321
0
    return NULL;
2322
0
}
2323
2324
static void
2325
combinations_dealloc(PyObject *op)
2326
0
{
2327
0
    combinationsobject *co = combinationsobject_CAST(op);
2328
0
    PyTypeObject *tp = Py_TYPE(co);
2329
0
    PyObject_GC_UnTrack(co);
2330
0
    Py_XDECREF(co->pool);
2331
0
    Py_XDECREF(co->result);
2332
0
    PyMem_Free(co->indices);
2333
0
    tp->tp_free(co);
2334
0
    Py_DECREF(tp);
2335
0
}
2336
2337
static PyObject *
2338
combinations_sizeof(PyObject *op, PyObject *Py_UNUSED(args))
2339
0
{
2340
0
    combinationsobject *co = combinationsobject_CAST(op);
2341
0
    size_t res = _PyObject_SIZE(Py_TYPE(co));
2342
0
    res += (size_t)co->r * sizeof(Py_ssize_t);
2343
0
    return PyLong_FromSize_t(res);
2344
0
}
2345
2346
static int
2347
combinations_traverse(PyObject *op, visitproc visit, void *arg)
2348
0
{
2349
0
    combinationsobject *co = combinationsobject_CAST(op);
2350
0
    Py_VISIT(Py_TYPE(co));
2351
0
    Py_VISIT(co->pool);
2352
0
    Py_VISIT(co->result);
2353
0
    return 0;
2354
0
}
2355
2356
static PyObject *
2357
combinations_next_lock_held(PyObject *op)
2358
0
{
2359
0
    combinationsobject *co = combinationsobject_CAST(op);
2360
0
    PyObject *elem;
2361
0
    PyObject *oldelem;
2362
0
    PyObject *pool = co->pool;
2363
0
    Py_ssize_t *indices = co->indices;
2364
0
    PyObject *result = co->result;
2365
0
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2366
0
    Py_ssize_t r = co->r;
2367
0
    Py_ssize_t i, j, index;
2368
2369
0
    if (co->stopped)
2370
0
        return NULL;
2371
2372
0
    if (result == NULL) {
2373
        /* On the first pass, initialize result tuple using the indices */
2374
0
        result = PyTuple_New(r);
2375
0
        if (result == NULL)
2376
0
            goto empty;
2377
0
        co->result = result;
2378
0
        for (i=0; i<r ; i++) {
2379
0
            index = indices[i];
2380
0
            elem = PyTuple_GET_ITEM(pool, index);
2381
0
            Py_INCREF(elem);
2382
0
            PyTuple_SET_ITEM(result, i, elem);
2383
0
        }
2384
0
    } else {
2385
        /* Copy the previous result tuple or re-use it if available */
2386
0
        if (!_PyObject_IsUniquelyReferenced(result)) {
2387
0
            PyObject *old_result = result;
2388
0
            result = PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2389
0
            if (result == NULL)
2390
0
                goto empty;
2391
0
            co->result = result;
2392
0
            Py_DECREF(old_result);
2393
0
        }
2394
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2395
        // recycling it, make sure it's tracked again:
2396
0
        else {
2397
0
            _PyTuple_Recycle(result);
2398
0
        }
2399
        /* Now, we've got the only copy so we can update it in-place
2400
         * CPython's empty tuple is a singleton and cached in
2401
         * PyTuple's freelist.
2402
         */
2403
0
        assert(r == 0 || Py_REFCNT(result) == 1);
2404
2405
        /* Scan indices right-to-left until finding one that is not
2406
           at its maximum (i + n - r). */
2407
0
        for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2408
0
            ;
2409
2410
        /* If i is negative, then the indices are all at
2411
           their maximum value and we're done. */
2412
0
        if (i < 0)
2413
0
            goto empty;
2414
2415
        /* Increment the current index which we know is not at its
2416
           maximum.  Then move back to the right setting each index
2417
           to its lowest possible value (one higher than the index
2418
           to its left -- this maintains the sort order invariant). */
2419
0
        indices[i]++;
2420
0
        for (j=i+1 ; j<r ; j++)
2421
0
            indices[j] = indices[j-1] + 1;
2422
2423
        /* Update the result tuple for the new indices
2424
           starting with i, the leftmost index that changed */
2425
0
        for ( ; i<r ; i++) {
2426
0
            index = indices[i];
2427
0
            elem = PyTuple_GET_ITEM(pool, index);
2428
0
            Py_INCREF(elem);
2429
0
            oldelem = PyTuple_GET_ITEM(result, i);
2430
0
            PyTuple_SET_ITEM(result, i, elem);
2431
0
            Py_DECREF(oldelem);
2432
0
        }
2433
0
    }
2434
2435
0
    return Py_NewRef(result);
2436
2437
0
empty:
2438
0
    co->stopped = 1;
2439
0
    return NULL;
2440
0
}
2441
2442
static PyObject *
2443
combinations_next(PyObject *op)
2444
0
{
2445
0
    PyObject *result;
2446
0
    Py_BEGIN_CRITICAL_SECTION(op);
2447
0
    result = combinations_next_lock_held(op);
2448
0
    Py_END_CRITICAL_SECTION()
2449
0
    return result;
2450
0
}
2451
2452
static PyMethodDef combinations_methods[] = {
2453
    {"__sizeof__", combinations_sizeof, METH_NOARGS, sizeof_doc},
2454
    {NULL,              NULL}   /* sentinel */
2455
};
2456
2457
static PyType_Slot combinations_slots[] = {
2458
    {Py_tp_dealloc, combinations_dealloc},
2459
    {Py_tp_getattro, PyObject_GenericGetAttr},
2460
    {Py_tp_doc, (void *)itertools_combinations__doc__},
2461
    {Py_tp_traverse, combinations_traverse},
2462
    {Py_tp_iter, PyObject_SelfIter},
2463
    {Py_tp_iternext, combinations_next},
2464
    {Py_tp_methods, combinations_methods},
2465
    {Py_tp_new, itertools_combinations},
2466
    {Py_tp_free, PyObject_GC_Del},
2467
    {0, NULL},
2468
};
2469
2470
static PyType_Spec combinations_spec = {
2471
    .name = "itertools.combinations",
2472
    .basicsize = sizeof(combinationsobject),
2473
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2474
              Py_TPFLAGS_IMMUTABLETYPE),
2475
    .slots = combinations_slots,
2476
};
2477
2478
2479
/* combinations with replacement object **************************************/
2480
2481
/* Equivalent to:
2482
2483
        def combinations_with_replacement(iterable, r):
2484
            "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2485
            # number items returned:  (n+r-1)! / r! / (n-1)!
2486
            pool = tuple(iterable)
2487
            n = len(pool)
2488
            indices = [0] * r
2489
            yield tuple(pool[i] for i in indices)
2490
            while 1:
2491
                for i in reversed(range(r)):
2492
                    if indices[i] != n - 1:
2493
                        break
2494
                else:
2495
                    return
2496
                indices[i:] = [indices[i] + 1] * (r - i)
2497
                yield tuple(pool[i] for i in indices)
2498
2499
        def combinations_with_replacement2(iterable, r):
2500
            'Alternate version that filters from product()'
2501
            pool = tuple(iterable)
2502
            n = len(pool)
2503
            for indices in product(range(n), repeat=r):
2504
                if sorted(indices) == list(indices):
2505
                    yield tuple(pool[i] for i in indices)
2506
*/
2507
typedef struct {
2508
    PyObject_HEAD
2509
    PyObject *pool;         /* input converted to a tuple */
2510
    Py_ssize_t *indices;    /* one index per result element */
2511
    PyObject *result;       /* most recently returned result tuple */
2512
    Py_ssize_t r;           /* size of result tuple */
2513
    int stopped;            /* set to 1 when the cwr iterator is exhausted */
2514
} cwrobject;
2515
2516
0
#define cwrobject_CAST(op)  ((cwrobject *)(op))
2517
2518
/*[clinic input]
2519
@permit_long_summary
2520
@permit_long_docstring_body
2521
@classmethod
2522
itertools.combinations_with_replacement.__new__
2523
    iterable: object
2524
    r: Py_ssize_t(allow_negative=False)
2525
Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2526
2527
combinations_with_replacement('ABC', 2) --> ('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')
2528
[clinic start generated code]*/
2529
2530
static PyObject *
2531
itertools_combinations_with_replacement_impl(PyTypeObject *type,
2532
                                             PyObject *iterable,
2533
                                             Py_ssize_t r)
2534
/*[clinic end generated code: output=48b26856d4e659ca input=828696750169e84f]*/
2535
0
{
2536
0
    cwrobject *co;
2537
0
    Py_ssize_t n;
2538
0
    PyObject *pool = NULL;
2539
0
    Py_ssize_t *indices = NULL;
2540
0
    Py_ssize_t i;
2541
2542
0
    pool = PySequence_Tuple(iterable);
2543
0
    if (pool == NULL)
2544
0
        goto error;
2545
0
    n = PyTuple_GET_SIZE(pool);
2546
2547
0
    indices = PyMem_New(Py_ssize_t, r);
2548
0
    if (indices == NULL) {
2549
0
        PyErr_NoMemory();
2550
0
        goto error;
2551
0
    }
2552
2553
0
    for (i=0 ; i<r ; i++)
2554
0
        indices[i] = 0;
2555
2556
    /* create cwrobject structure */
2557
0
    co = (cwrobject *)type->tp_alloc(type, 0);
2558
0
    if (co == NULL)
2559
0
        goto error;
2560
2561
0
    co->pool = pool;
2562
0
    co->indices = indices;
2563
0
    co->result = NULL;
2564
0
    co->r = r;
2565
0
    co->stopped = !n && r;
2566
2567
0
    return (PyObject *)co;
2568
2569
0
error:
2570
0
    if (indices != NULL)
2571
0
        PyMem_Free(indices);
2572
0
    Py_XDECREF(pool);
2573
0
    return NULL;
2574
0
}
2575
2576
static void
2577
cwr_dealloc(PyObject *op)
2578
0
{
2579
0
    cwrobject *co = cwrobject_CAST(op);
2580
0
    PyTypeObject *tp = Py_TYPE(co);
2581
0
    PyObject_GC_UnTrack(co);
2582
0
    Py_XDECREF(co->pool);
2583
0
    Py_XDECREF(co->result);
2584
0
    PyMem_Free(co->indices);
2585
0
    tp->tp_free(co);
2586
0
    Py_DECREF(tp);
2587
0
}
2588
2589
static PyObject *
2590
cwr_sizeof(PyObject *op, PyObject *Py_UNUSED(args))
2591
0
{
2592
0
    cwrobject *co = cwrobject_CAST(op);
2593
0
    size_t res = _PyObject_SIZE(Py_TYPE(co));
2594
0
    res += (size_t)co->r * sizeof(Py_ssize_t);
2595
0
    return PyLong_FromSize_t(res);
2596
0
}
2597
2598
static int
2599
cwr_traverse(PyObject *op, visitproc visit, void *arg)
2600
0
{
2601
0
    cwrobject *co = cwrobject_CAST(op);
2602
0
    Py_VISIT(Py_TYPE(co));
2603
0
    Py_VISIT(co->pool);
2604
0
    Py_VISIT(co->result);
2605
0
    return 0;
2606
0
}
2607
2608
static PyObject *
2609
cwr_next_lock_held(PyObject *op)
2610
0
{
2611
0
    cwrobject *co = cwrobject_CAST(op);
2612
0
    PyObject *elem;
2613
0
    PyObject *oldelem;
2614
0
    PyObject *pool = co->pool;
2615
0
    Py_ssize_t *indices = co->indices;
2616
0
    PyObject *result = co->result;
2617
0
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2618
0
    Py_ssize_t r = co->r;
2619
0
    Py_ssize_t i, index;
2620
2621
0
    if (co->stopped)
2622
0
        return NULL;
2623
2624
0
    if (result == NULL) {
2625
        /* On the first pass, initialize result tuple with pool[0] */
2626
0
        result = PyTuple_New(r);
2627
0
        if (result == NULL)
2628
0
            goto empty;
2629
0
        co->result = result;
2630
0
        if (n > 0) {
2631
0
            elem = PyTuple_GET_ITEM(pool, 0);
2632
0
            for (i=0; i<r ; i++) {
2633
0
                assert(indices[i] == 0);
2634
0
                Py_INCREF(elem);
2635
0
                PyTuple_SET_ITEM(result, i, elem);
2636
0
            }
2637
0
        }
2638
0
    } else {
2639
        /* Copy the previous result tuple or re-use it if available */
2640
0
        if (!_PyObject_IsUniquelyReferenced(result)) {
2641
0
            PyObject *old_result = result;
2642
0
            result = PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2643
0
            if (result == NULL)
2644
0
                goto empty;
2645
0
            co->result = result;
2646
0
            Py_DECREF(old_result);
2647
0
        }
2648
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2649
        // recycling it, make sure it's tracked again:
2650
0
        else {
2651
0
            _PyTuple_Recycle(result);
2652
0
        }
2653
        /* Now, we've got the only copy so we can update it in-place CPython's
2654
           empty tuple is a singleton and cached in PyTuple's freelist. */
2655
0
        assert(r == 0 || Py_REFCNT(result) == 1);
2656
2657
       /* Scan indices right-to-left until finding one that is not
2658
        * at its maximum (n-1). */
2659
0
        for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2660
0
            ;
2661
2662
        /* If i is negative, then the indices are all at
2663
           their maximum value and we're done. */
2664
0
        if (i < 0)
2665
0
            goto empty;
2666
2667
        /* Increment the current index which we know is not at its
2668
           maximum.  Then set all to the right to the same value. */
2669
0
        index = indices[i] + 1;
2670
0
        assert(index < n);
2671
0
        elem = PyTuple_GET_ITEM(pool, index);
2672
0
        for ( ; i<r ; i++) {
2673
0
            indices[i] = index;
2674
0
            Py_INCREF(elem);
2675
0
            oldelem = PyTuple_GET_ITEM(result, i);
2676
0
            PyTuple_SET_ITEM(result, i, elem);
2677
0
            Py_DECREF(oldelem);
2678
0
        }
2679
0
    }
2680
2681
0
    return Py_NewRef(result);
2682
2683
0
empty:
2684
0
    co->stopped = 1;
2685
0
    return NULL;
2686
0
}
2687
2688
static PyObject *
2689
cwr_next(PyObject *op)
2690
0
{
2691
0
    PyObject *result;
2692
0
    Py_BEGIN_CRITICAL_SECTION(op);
2693
0
    result = cwr_next_lock_held(op);
2694
0
    Py_END_CRITICAL_SECTION()
2695
0
    return result;
2696
0
}
2697
2698
static PyMethodDef cwr_methods[] = {
2699
    {"__sizeof__", cwr_sizeof, METH_NOARGS, sizeof_doc},
2700
    {NULL,              NULL}   /* sentinel */
2701
};
2702
2703
static PyType_Slot cwr_slots[] = {
2704
    {Py_tp_dealloc, cwr_dealloc},
2705
    {Py_tp_getattro, PyObject_GenericGetAttr},
2706
    {Py_tp_doc, (void *)itertools_combinations_with_replacement__doc__},
2707
    {Py_tp_traverse, cwr_traverse},
2708
    {Py_tp_iter, PyObject_SelfIter},
2709
    {Py_tp_iternext, cwr_next},
2710
    {Py_tp_methods, cwr_methods},
2711
    {Py_tp_new, itertools_combinations_with_replacement},
2712
    {Py_tp_free, PyObject_GC_Del},
2713
    {0, NULL},
2714
};
2715
2716
static PyType_Spec cwr_spec = {
2717
    .name = "itertools.combinations_with_replacement",
2718
    .basicsize = sizeof(cwrobject),
2719
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2720
              Py_TPFLAGS_IMMUTABLETYPE),
2721
    .slots = cwr_slots,
2722
};
2723
2724
2725
/* permutations object ********************************************************
2726
2727
def permutations(iterable, r=None):
2728
    # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
2729
    # permutations(range(3)) --> 012 021 102 120 201 210
2730
    pool = tuple(iterable)
2731
    n = len(pool)
2732
    r = n if r is None else r
2733
    if r > n:
2734
        return
2735
    indices = list(range(n))
2736
    cycles = list(range(n, n-r, -1))
2737
    yield tuple(pool[i] for i in indices[:r])
2738
    while n:
2739
        for i in reversed(range(r)):
2740
            cycles[i] -= 1
2741
            if cycles[i] == 0:
2742
                indices[i:] = indices[i+1:] + indices[i:i+1]
2743
                cycles[i] = n - i
2744
            else:
2745
                j = cycles[i]
2746
                indices[i], indices[-j] = indices[-j], indices[i]
2747
                yield tuple(pool[i] for i in indices[:r])
2748
                break
2749
        else:
2750
            return
2751
*/
2752
2753
typedef struct {
2754
    PyObject_HEAD
2755
    PyObject *pool;         /* input converted to a tuple */
2756
    Py_ssize_t *indices;    /* one index per element in the pool */
2757
    Py_ssize_t *cycles;     /* one rollover counter per element in the result */
2758
    PyObject *result;       /* most recently returned result tuple */
2759
    Py_ssize_t r;           /* size of result tuple */
2760
    int stopped;            /* set to 1 when the iterator is exhausted */
2761
} permutationsobject;
2762
2763
162
#define permutationsobject_CAST(op) ((permutationsobject *)(op))
2764
2765
/*[clinic input]
2766
@classmethod
2767
itertools.permutations.__new__
2768
    iterable: object
2769
    r as robj: object = None
2770
Return successive r-length permutations of elements in the iterable.
2771
2772
permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
2773
[clinic start generated code]*/
2774
2775
static PyObject *
2776
itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
2777
                            PyObject *robj)
2778
/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
2779
48
{
2780
48
    permutationsobject *po;
2781
48
    Py_ssize_t n;
2782
48
    Py_ssize_t r;
2783
48
    PyObject *pool = NULL;
2784
48
    Py_ssize_t *indices = NULL;
2785
48
    Py_ssize_t *cycles = NULL;
2786
48
    Py_ssize_t i;
2787
2788
48
    pool = PySequence_Tuple(iterable);
2789
48
    if (pool == NULL)
2790
0
        goto error;
2791
48
    n = PyTuple_GET_SIZE(pool);
2792
2793
48
    r = n;
2794
48
    if (robj != Py_None) {
2795
0
        if (!PyLong_Check(robj)) {
2796
0
            PyErr_SetString(PyExc_TypeError, "Expected int as r");
2797
0
            goto error;
2798
0
        }
2799
0
        r = PyLong_AsSsize_t(robj);
2800
0
        if (r == -1 && PyErr_Occurred())
2801
0
            goto error;
2802
0
    }
2803
48
    if (r < 0) {
2804
0
        PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2805
0
        goto error;
2806
0
    }
2807
2808
48
    indices = PyMem_New(Py_ssize_t, n);
2809
48
    cycles = PyMem_New(Py_ssize_t, r);
2810
48
    if (indices == NULL || cycles == NULL) {
2811
0
        PyErr_NoMemory();
2812
0
        goto error;
2813
0
    }
2814
2815
114
    for (i=0 ; i<n ; i++)
2816
66
        indices[i] = i;
2817
114
    for (i=0 ; i<r ; i++)
2818
66
        cycles[i] = n - i;
2819
2820
    /* create permutationsobject structure */
2821
48
    po = (permutationsobject *)type->tp_alloc(type, 0);
2822
48
    if (po == NULL)
2823
0
        goto error;
2824
2825
48
    po->pool = pool;
2826
48
    po->indices = indices;
2827
48
    po->cycles = cycles;
2828
48
    po->result = NULL;
2829
48
    po->r = r;
2830
48
    po->stopped = r > n ? 1 : 0;
2831
2832
48
    return (PyObject *)po;
2833
2834
0
error:
2835
0
    if (indices != NULL)
2836
0
        PyMem_Free(indices);
2837
0
    if (cycles != NULL)
2838
0
        PyMem_Free(cycles);
2839
0
    Py_XDECREF(pool);
2840
0
    return NULL;
2841
48
}
2842
2843
static void
2844
permutations_dealloc(PyObject *op)
2845
48
{
2846
48
    permutationsobject *po = permutationsobject_CAST(op);
2847
48
    PyTypeObject *tp = Py_TYPE(po);
2848
48
    PyObject_GC_UnTrack(po);
2849
48
    Py_XDECREF(po->pool);
2850
48
    Py_XDECREF(po->result);
2851
48
    PyMem_Free(po->indices);
2852
48
    PyMem_Free(po->cycles);
2853
48
    tp->tp_free(po);
2854
48
    Py_DECREF(tp);
2855
48
}
2856
2857
static PyObject *
2858
permutations_sizeof(PyObject *op, PyObject *Py_UNUSED(args))
2859
0
{
2860
0
    permutationsobject *po = permutationsobject_CAST(op);
2861
0
    size_t res = _PyObject_SIZE(Py_TYPE(po));
2862
0
    res += (size_t)PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
2863
0
    res += (size_t)po->r * sizeof(Py_ssize_t);
2864
0
    return PyLong_FromSize_t(res);
2865
0
}
2866
2867
static int
2868
permutations_traverse(PyObject *op, visitproc visit, void *arg)
2869
0
{
2870
0
    permutationsobject *po = permutationsobject_CAST(op);
2871
0
    Py_VISIT(Py_TYPE(po));
2872
0
    Py_VISIT(po->pool);
2873
0
    Py_VISIT(po->result);
2874
0
    return 0;
2875
0
}
2876
2877
static PyObject *
2878
permutations_next_lock_held(PyObject *op)
2879
114
{
2880
114
    permutationsobject *po = permutationsobject_CAST(op);
2881
114
    PyObject *elem;
2882
114
    PyObject *oldelem;
2883
114
    PyObject *pool = po->pool;
2884
114
    Py_ssize_t *indices = po->indices;
2885
114
    Py_ssize_t *cycles = po->cycles;
2886
114
    PyObject *result = po->result;
2887
114
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2888
114
    Py_ssize_t r = po->r;
2889
114
    Py_ssize_t i, j, k, index;
2890
2891
114
    if (po->stopped)
2892
0
        return NULL;
2893
2894
114
    if (result == NULL) {
2895
        /* On the first pass, initialize result tuple using the indices */
2896
48
        result = PyTuple_New(r);
2897
48
        if (result == NULL)
2898
0
            goto empty;
2899
48
        po->result = result;
2900
114
        for (i=0; i<r ; i++) {
2901
66
            index = indices[i];
2902
66
            elem = PyTuple_GET_ITEM(pool, index);
2903
66
            Py_INCREF(elem);
2904
66
            PyTuple_SET_ITEM(result, i, elem);
2905
66
        }
2906
66
    } else {
2907
66
        if (n == 0)
2908
0
            goto empty;
2909
2910
        /* Copy the previous result tuple or re-use it if available */
2911
66
        if (!_PyObject_IsUniquelyReferenced(result)) {
2912
66
            PyObject *old_result = result;
2913
66
            result = PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2914
66
            if (result == NULL)
2915
0
                goto empty;
2916
66
            po->result = result;
2917
66
            Py_DECREF(old_result);
2918
66
        }
2919
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2920
        // recycling it, make sure it's tracked again:
2921
0
        else {
2922
0
            _PyTuple_Recycle(result);
2923
0
        }
2924
        /* Now, we've got the only copy so we can update it in-place */
2925
66
        assert(r == 0 || Py_REFCNT(result) == 1);
2926
2927
        /* Decrement rightmost cycle, moving leftward upon zero rollover */
2928
150
        for (i=r-1 ; i>=0 ; i--) {
2929
102
            cycles[i] -= 1;
2930
102
            if (cycles[i] == 0) {
2931
                /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2932
84
                index = indices[i];
2933
102
                for (j=i ; j<n-1 ; j++)
2934
18
                    indices[j] = indices[j+1];
2935
84
                indices[n-1] = index;
2936
84
                cycles[i] = n - i;
2937
84
            } else {
2938
18
                j = cycles[i];
2939
18
                index = indices[i];
2940
18
                indices[i] = indices[n-j];
2941
18
                indices[n-j] = index;
2942
2943
54
                for (k=i; k<r ; k++) {
2944
                    /* start with i, the leftmost element that changed */
2945
                    /* yield tuple(pool[k] for k in indices[:r]) */
2946
36
                    index = indices[k];
2947
36
                    elem = PyTuple_GET_ITEM(pool, index);
2948
36
                    Py_INCREF(elem);
2949
36
                    oldelem = PyTuple_GET_ITEM(result, k);
2950
36
                    PyTuple_SET_ITEM(result, k, elem);
2951
36
                    Py_DECREF(oldelem);
2952
36
                }
2953
18
                break;
2954
18
            }
2955
102
        }
2956
        /* If i is negative, then the cycles have all
2957
           rolled-over and we're done. */
2958
66
        if (i < 0)
2959
48
            goto empty;
2960
66
    }
2961
66
    return Py_NewRef(result);
2962
2963
48
empty:
2964
48
    po->stopped = 1;
2965
48
    return NULL;
2966
114
}
2967
2968
static PyObject *
2969
permutations_next(PyObject *op)
2970
114
{
2971
114
    PyObject *result;
2972
114
    Py_BEGIN_CRITICAL_SECTION(op);
2973
114
    result = permutations_next_lock_held(op);
2974
114
    Py_END_CRITICAL_SECTION()
2975
114
    return result;
2976
114
}
2977
2978
static PyMethodDef permuations_methods[] = {
2979
    {"__sizeof__", permutations_sizeof, METH_NOARGS, sizeof_doc},
2980
    {NULL,              NULL}   /* sentinel */
2981
};
2982
2983
static PyType_Slot permutations_slots[] = {
2984
    {Py_tp_dealloc, permutations_dealloc},
2985
    {Py_tp_getattro, PyObject_GenericGetAttr},
2986
    {Py_tp_doc, (void *)itertools_permutations__doc__},
2987
    {Py_tp_traverse, permutations_traverse},
2988
    {Py_tp_iter, PyObject_SelfIter},
2989
    {Py_tp_iternext, permutations_next},
2990
    {Py_tp_methods, permuations_methods},
2991
    {Py_tp_new, itertools_permutations},
2992
    {Py_tp_free, PyObject_GC_Del},
2993
    {0, NULL},
2994
};
2995
2996
static PyType_Spec permutations_spec = {
2997
    .name = "itertools.permutations",
2998
    .basicsize = sizeof(permutationsobject),
2999
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3000
              Py_TPFLAGS_IMMUTABLETYPE),
3001
    .slots = permutations_slots,
3002
};
3003
3004
3005
/* accumulate object ********************************************************/
3006
3007
typedef struct {
3008
    PyObject_HEAD
3009
    PyObject *total;
3010
    PyObject *it;
3011
    PyObject *binop;
3012
    PyObject *initial;
3013
    itertools_state *state;
3014
} accumulateobject;
3015
3016
0
#define accumulateobject_CAST(op)   ((accumulateobject *)(op))
3017
3018
/*[clinic input]
3019
@classmethod
3020
itertools.accumulate.__new__
3021
    iterable: object
3022
    func as binop: object = None
3023
    *
3024
    initial: object = None
3025
Return series of accumulated sums (or other binary function results).
3026
[clinic start generated code]*/
3027
3028
static PyObject *
3029
itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
3030
                          PyObject *binop, PyObject *initial)
3031
/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
3032
0
{
3033
0
    PyObject *it;
3034
0
    accumulateobject *lz;
3035
3036
    /* Get iterator. */
3037
0
    it = PyObject_GetIter(iterable);
3038
0
    if (it == NULL)
3039
0
        return NULL;
3040
3041
    /* create accumulateobject structure */
3042
0
    lz = (accumulateobject *)type->tp_alloc(type, 0);
3043
0
    if (lz == NULL) {
3044
0
        Py_DECREF(it);
3045
0
        return NULL;
3046
0
    }
3047
3048
0
    if (binop != Py_None) {
3049
0
        lz->binop = Py_XNewRef(binop);
3050
0
    }
3051
0
    lz->total = NULL;
3052
0
    lz->it = it;
3053
0
    lz->initial = Py_XNewRef(initial);
3054
0
    lz->state = find_state_by_type(type);
3055
0
    return (PyObject *)lz;
3056
0
}
3057
3058
static void
3059
accumulate_dealloc(PyObject *op)
3060
0
{
3061
0
    accumulateobject *lz = accumulateobject_CAST(op);
3062
0
    PyTypeObject *tp = Py_TYPE(lz);
3063
0
    PyObject_GC_UnTrack(lz);
3064
0
    Py_XDECREF(lz->binop);
3065
0
    Py_XDECREF(lz->total);
3066
0
    Py_XDECREF(lz->it);
3067
0
    Py_XDECREF(lz->initial);
3068
0
    tp->tp_free(lz);
3069
0
    Py_DECREF(tp);
3070
0
}
3071
3072
static int
3073
accumulate_traverse(PyObject *op, visitproc visit, void *arg)
3074
0
{
3075
0
    accumulateobject *lz = accumulateobject_CAST(op);
3076
0
    Py_VISIT(Py_TYPE(lz));
3077
0
    Py_VISIT(lz->binop);
3078
0
    Py_VISIT(lz->it);
3079
0
    Py_VISIT(lz->total);
3080
0
    Py_VISIT(lz->initial);
3081
0
    return 0;
3082
0
}
3083
3084
static PyObject *
3085
accumulate_next_lock_held(PyObject *op)
3086
0
{
3087
0
    accumulateobject *lz = accumulateobject_CAST(op);
3088
0
    PyObject *val, *newtotal;
3089
3090
0
    if (lz->initial != Py_None) {
3091
0
        lz->total = lz->initial;
3092
0
        lz->initial = Py_NewRef(Py_None);
3093
0
        return Py_NewRef(lz->total);
3094
0
    }
3095
0
    val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
3096
0
    if (val == NULL)
3097
0
        return NULL;
3098
3099
0
    if (lz->total == NULL) {
3100
0
        lz->total = Py_NewRef(val);
3101
0
        return lz->total;
3102
0
    }
3103
3104
0
    if (lz->binop == NULL)
3105
0
        newtotal = PyNumber_Add(lz->total, val);
3106
0
    else
3107
0
        newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
3108
0
    Py_DECREF(val);
3109
0
    if (newtotal == NULL)
3110
0
        return NULL;
3111
3112
0
    Py_INCREF(newtotal);
3113
0
    Py_SETREF(lz->total, newtotal);
3114
0
    return newtotal;
3115
0
}
3116
3117
static PyObject *
3118
accumulate_next(PyObject *op)
3119
0
{
3120
0
    PyObject *result;
3121
0
    Py_BEGIN_CRITICAL_SECTION(op);
3122
0
    result = accumulate_next_lock_held(op);
3123
0
    Py_END_CRITICAL_SECTION()
3124
0
    return result;
3125
0
}
3126
3127
static PyType_Slot accumulate_slots[] = {
3128
    {Py_tp_dealloc, accumulate_dealloc},
3129
    {Py_tp_getattro, PyObject_GenericGetAttr},
3130
    {Py_tp_doc, (void *)itertools_accumulate__doc__},
3131
    {Py_tp_traverse, accumulate_traverse},
3132
    {Py_tp_iter, PyObject_SelfIter},
3133
    {Py_tp_iternext, accumulate_next},
3134
    {Py_tp_new, itertools_accumulate},
3135
    {Py_tp_free, PyObject_GC_Del},
3136
    {0, NULL},
3137
};
3138
3139
static PyType_Spec accumulate_spec = {
3140
    .name = "itertools.accumulate",
3141
    .basicsize = sizeof(accumulateobject),
3142
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3143
              Py_TPFLAGS_IMMUTABLETYPE),
3144
    .slots = accumulate_slots,
3145
};
3146
3147
3148
/* compress object ************************************************************/
3149
3150
/* Equivalent to:
3151
3152
    def compress(data, selectors):
3153
        "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3154
        return (d for d, s in zip(data, selectors) if s)
3155
*/
3156
3157
typedef struct {
3158
    PyObject_HEAD
3159
    PyObject *data;
3160
    PyObject *selectors;
3161
} compressobject;
3162
3163
0
#define compressobject_CAST(op) ((compressobject *)(op))
3164
3165
/*[clinic input]
3166
@classmethod
3167
itertools.compress.__new__
3168
    data as seq1: object
3169
    selectors as seq2: object
3170
Return data elements corresponding to true selector elements.
3171
3172
Forms a shorter iterator from selected data elements using the selectors
3173
to choose the data elements.
3174
[clinic start generated code]*/
3175
3176
static PyObject *
3177
itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3178
/*[clinic end generated code: output=7e67157212ed09e0 input=32ca4347dbc46749]*/
3179
0
{
3180
0
    PyObject *data=NULL, *selectors=NULL;
3181
0
    compressobject *lz;
3182
3183
0
    data = PyObject_GetIter(seq1);
3184
0
    if (data == NULL)
3185
0
        goto fail;
3186
0
    selectors = PyObject_GetIter(seq2);
3187
0
    if (selectors == NULL)
3188
0
        goto fail;
3189
3190
    /* create compressobject structure */
3191
0
    lz = (compressobject *)type->tp_alloc(type, 0);
3192
0
    if (lz == NULL)
3193
0
        goto fail;
3194
0
    lz->data = data;
3195
0
    lz->selectors = selectors;
3196
0
    return (PyObject *)lz;
3197
3198
0
fail:
3199
0
    Py_XDECREF(data);
3200
0
    Py_XDECREF(selectors);
3201
0
    return NULL;
3202
0
}
3203
3204
static void
3205
compress_dealloc(PyObject *op)
3206
0
{
3207
0
    compressobject *lz = compressobject_CAST(op);
3208
0
    PyTypeObject *tp = Py_TYPE(lz);
3209
0
    PyObject_GC_UnTrack(lz);
3210
0
    Py_XDECREF(lz->data);
3211
0
    Py_XDECREF(lz->selectors);
3212
0
    tp->tp_free(lz);
3213
0
    Py_DECREF(tp);
3214
0
}
3215
3216
static int
3217
compress_traverse(PyObject *op, visitproc visit, void *arg)
3218
0
{
3219
0
    compressobject *lz = compressobject_CAST(op);
3220
0
    Py_VISIT(Py_TYPE(lz));
3221
0
    Py_VISIT(lz->data);
3222
0
    Py_VISIT(lz->selectors);
3223
0
    return 0;
3224
0
}
3225
3226
static PyObject *
3227
compress_next(PyObject *op)
3228
0
{
3229
0
    compressobject *lz = compressobject_CAST(op);
3230
0
    PyObject *data = lz->data, *selectors = lz->selectors;
3231
0
    PyObject *datum, *selector;
3232
0
    PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3233
0
    PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3234
0
    int ok;
3235
3236
0
    while (1) {
3237
        /* Steps:  get datum, get selector, evaluate selector.
3238
           Order is important (to match the pure python version
3239
           in terms of which input gets a chance to raise an
3240
           exception first).
3241
        */
3242
3243
0
        datum = datanext(data);
3244
0
        if (datum == NULL)
3245
0
            return NULL;
3246
3247
0
        selector = selectornext(selectors);
3248
0
        if (selector == NULL) {
3249
0
            Py_DECREF(datum);
3250
0
            return NULL;
3251
0
        }
3252
3253
0
        ok = PyObject_IsTrue(selector);
3254
0
        Py_DECREF(selector);
3255
0
        if (ok > 0)
3256
0
            return datum;
3257
0
        Py_DECREF(datum);
3258
0
        if (ok < 0)
3259
0
            return NULL;
3260
0
    }
3261
0
}
3262
3263
static PyType_Slot compress_slots[] = {
3264
    {Py_tp_dealloc, compress_dealloc},
3265
    {Py_tp_getattro, PyObject_GenericGetAttr},
3266
    {Py_tp_doc, (void *)itertools_compress__doc__},
3267
    {Py_tp_traverse, compress_traverse},
3268
    {Py_tp_iter, PyObject_SelfIter},
3269
    {Py_tp_iternext, compress_next},
3270
    {Py_tp_new, itertools_compress},
3271
    {Py_tp_free, PyObject_GC_Del},
3272
    {0, NULL},
3273
};
3274
3275
static PyType_Spec compress_spec = {
3276
    .name = "itertools.compress",
3277
    .basicsize = sizeof(compressobject),
3278
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3279
              Py_TPFLAGS_IMMUTABLETYPE),
3280
    .slots = compress_slots,
3281
};
3282
3283
3284
/* filterfalse object ************************************************************/
3285
3286
typedef struct {
3287
    PyObject_HEAD
3288
    PyObject *func;
3289
    PyObject *it;
3290
} filterfalseobject;
3291
3292
0
#define filterfalseobject_CAST(op)  ((filterfalseobject *)(op))
3293
3294
/*[clinic input]
3295
@classmethod
3296
itertools.filterfalse.__new__
3297
    function as func: object
3298
    iterable as seq: object
3299
    /
3300
Return those items of iterable for which function(item) is false.
3301
3302
If function is None, return the items that are false.
3303
[clinic start generated code]*/
3304
3305
static PyObject *
3306
itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3307
/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
3308
0
{
3309
0
    PyObject *it;
3310
0
    filterfalseobject *lz;
3311
3312
    /* Get iterator. */
3313
0
    it = PyObject_GetIter(seq);
3314
0
    if (it == NULL)
3315
0
        return NULL;
3316
3317
    /* create filterfalseobject structure */
3318
0
    lz = (filterfalseobject *)type->tp_alloc(type, 0);
3319
0
    if (lz == NULL) {
3320
0
        Py_DECREF(it);
3321
0
        return NULL;
3322
0
    }
3323
0
    lz->func = Py_NewRef(func);
3324
0
    lz->it = it;
3325
3326
0
    return (PyObject *)lz;
3327
0
}
3328
3329
static void
3330
filterfalse_dealloc(PyObject *op)
3331
0
{
3332
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3333
0
    PyTypeObject *tp = Py_TYPE(lz);
3334
0
    PyObject_GC_UnTrack(lz);
3335
0
    Py_XDECREF(lz->func);
3336
0
    Py_XDECREF(lz->it);
3337
0
    tp->tp_free(lz);
3338
0
    Py_DECREF(tp);
3339
0
}
3340
3341
static int
3342
filterfalse_traverse(PyObject *op, visitproc visit, void *arg)
3343
0
{
3344
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3345
0
    Py_VISIT(Py_TYPE(lz));
3346
0
    Py_VISIT(lz->it);
3347
0
    Py_VISIT(lz->func);
3348
0
    return 0;
3349
0
}
3350
3351
static PyObject *
3352
filterfalse_next(PyObject *op)
3353
0
{
3354
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3355
0
    PyObject *item;
3356
0
    PyObject *it = lz->it;
3357
0
    long ok;
3358
0
    PyObject *(*iternext)(PyObject *);
3359
3360
0
    iternext = *Py_TYPE(it)->tp_iternext;
3361
0
    for (;;) {
3362
0
        item = iternext(it);
3363
0
        if (item == NULL)
3364
0
            return NULL;
3365
3366
0
        if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3367
0
            ok = PyObject_IsTrue(item);
3368
0
        } else {
3369
0
            PyObject *good;
3370
0
            good = PyObject_CallOneArg(lz->func, item);
3371
0
            if (good == NULL) {
3372
0
                Py_DECREF(item);
3373
0
                return NULL;
3374
0
            }
3375
0
            ok = PyObject_IsTrue(good);
3376
0
            Py_DECREF(good);
3377
0
        }
3378
0
        if (ok == 0)
3379
0
            return item;
3380
0
        Py_DECREF(item);
3381
0
        if (ok < 0)
3382
0
            return NULL;
3383
0
    }
3384
0
}
3385
3386
static PyType_Slot filterfalse_slots[] = {
3387
    {Py_tp_dealloc, filterfalse_dealloc},
3388
    {Py_tp_getattro, PyObject_GenericGetAttr},
3389
    {Py_tp_doc, (void *)itertools_filterfalse__doc__},
3390
    {Py_tp_traverse, filterfalse_traverse},
3391
    {Py_tp_iter, PyObject_SelfIter},
3392
    {Py_tp_iternext, filterfalse_next},
3393
    {Py_tp_new, itertools_filterfalse},
3394
    {Py_tp_free, PyObject_GC_Del},
3395
    {0, NULL},
3396
};
3397
3398
static PyType_Spec filterfalse_spec = {
3399
    .name = "itertools.filterfalse",
3400
    .basicsize = sizeof(filterfalseobject),
3401
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3402
              Py_TPFLAGS_IMMUTABLETYPE),
3403
    .slots = filterfalse_slots,
3404
};
3405
3406
3407
/* count object ************************************************************/
3408
3409
typedef struct {
3410
    PyObject_HEAD
3411
    Py_ssize_t cnt;
3412
    PyObject *long_cnt;
3413
    PyObject *long_step;
3414
} countobject;
3415
3416
0
#define countobject_CAST(op)    ((countobject *)(op))
3417
3418
/* Counting logic and invariants:
3419
3420
fast_mode:  when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3421
3422
    assert(long_cnt == NULL && long_step==PyLong(1));
3423
    Advances with:  cnt += 1
3424
    When count hits PY_SSIZE_T_MAX, switch to slow_mode.
3425
3426
slow_mode:  when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3427
3428
    assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3429
    All counting is done with python objects (no overflows or underflows).
3430
    Advances with:  long_cnt += long_step
3431
    Step may be zero -- effectively a slow version of repeat(cnt).
3432
    Either long_cnt or long_step may be a float, Fraction, or Decimal.
3433
*/
3434
3435
/*[clinic input]
3436
@permit_long_summary
3437
@classmethod
3438
itertools.count.__new__
3439
    start as long_cnt: object(c_default="NULL") = 0
3440
    step as long_step: object(c_default="NULL") = 1
3441
Return a count object whose .__next__() method returns consecutive values.
3442
3443
Equivalent to:
3444
    def count(firstval=0, step=1):
3445
        x = firstval
3446
        while 1:
3447
            yield x
3448
            x += step
3449
[clinic start generated code]*/
3450
3451
static PyObject *
3452
itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
3453
                     PyObject *long_step)
3454
/*[clinic end generated code: output=09a9250aebd00b1c input=91e4b12c0e88b9f4]*/
3455
0
{
3456
0
    countobject *lz;
3457
0
    int fast_mode;
3458
0
    Py_ssize_t cnt = 0;
3459
0
    long step;
3460
3461
0
    if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3462
0
        (long_step != NULL && !PyNumber_Check(long_step))) {
3463
0
                    PyErr_SetString(PyExc_TypeError, "a number is required");
3464
0
                    return NULL;
3465
0
    }
3466
3467
0
    fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3468
0
                (long_step == NULL || PyLong_Check(long_step));
3469
3470
    /* If not specified, start defaults to 0 */
3471
0
    if (long_cnt != NULL) {
3472
0
        if (fast_mode) {
3473
0
            assert(PyLong_Check(long_cnt));
3474
0
            cnt = PyLong_AsSsize_t(long_cnt);
3475
0
            if (cnt == -1 && PyErr_Occurred()) {
3476
0
                PyErr_Clear();
3477
0
                fast_mode = 0;
3478
0
            }
3479
0
        }
3480
0
    } else {
3481
0
        cnt = 0;
3482
0
        long_cnt = _PyLong_GetZero();
3483
0
    }
3484
0
    Py_INCREF(long_cnt);
3485
3486
    /* If not specified, step defaults to 1 */
3487
0
    if (long_step == NULL) {
3488
0
        long_step = _PyLong_GetOne();
3489
0
    }
3490
0
    Py_INCREF(long_step);
3491
3492
0
    assert(long_cnt != NULL && long_step != NULL);
3493
3494
    /* Fast mode only works when the step is 1 */
3495
0
    if (fast_mode) {
3496
0
        assert(PyLong_Check(long_step));
3497
0
        step = PyLong_AsLong(long_step);
3498
0
        if (step != 1) {
3499
0
            fast_mode = 0;
3500
0
            if (step == -1 && PyErr_Occurred())
3501
0
                PyErr_Clear();
3502
0
        }
3503
0
    }
3504
3505
0
    if (fast_mode)
3506
0
        Py_CLEAR(long_cnt);
3507
0
    else
3508
0
        cnt = PY_SSIZE_T_MAX;
3509
3510
0
    assert((long_cnt == NULL && fast_mode) ||
3511
0
           (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
3512
0
    assert(!fast_mode ||
3513
0
           (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
3514
3515
    /* create countobject structure */
3516
0
    lz = (countobject *)type->tp_alloc(type, 0);
3517
0
    if (lz == NULL) {
3518
0
        Py_XDECREF(long_cnt);
3519
0
        Py_DECREF(long_step);
3520
0
        return NULL;
3521
0
    }
3522
0
    lz->cnt = cnt;
3523
0
    lz->long_cnt = long_cnt;
3524
0
    lz->long_step = long_step;
3525
3526
0
    return (PyObject *)lz;
3527
0
}
3528
3529
static void
3530
count_dealloc(PyObject *op)
3531
0
{
3532
0
    countobject *lz = countobject_CAST(op);
3533
0
    PyTypeObject *tp = Py_TYPE(lz);
3534
0
    PyObject_GC_UnTrack(lz);
3535
0
    Py_XDECREF(lz->long_cnt);
3536
0
    Py_XDECREF(lz->long_step);
3537
0
    tp->tp_free(lz);
3538
0
    Py_DECREF(tp);
3539
0
}
3540
3541
static int
3542
count_traverse(PyObject *op, visitproc visit, void *arg)
3543
0
{
3544
0
    countobject *lz = countobject_CAST(op);
3545
0
    Py_VISIT(Py_TYPE(lz));
3546
0
    Py_VISIT(lz->long_cnt);
3547
0
    Py_VISIT(lz->long_step);
3548
0
    return 0;
3549
0
}
3550
3551
static PyObject *
3552
count_nextlong(countobject *lz)
3553
0
{
3554
0
    if (lz->long_cnt == NULL) {
3555
        /* Switch to slow_mode */
3556
0
        lz->long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3557
0
        if (lz->long_cnt == NULL) {
3558
0
            return NULL;
3559
0
        }
3560
0
    }
3561
0
    assert(lz->cnt == PY_SSIZE_T_MAX && lz->long_cnt != NULL);
3562
3563
    // We hold one reference to "result" (a.k.a. the old value of
3564
    // lz->long_cnt); we'll either return it or keep it in lz->long_cnt.
3565
0
    PyObject *result = lz->long_cnt;
3566
3567
0
    PyObject *stepped_up = PyNumber_Add(result, lz->long_step);
3568
0
    if (stepped_up == NULL) {
3569
0
        return NULL;
3570
0
    }
3571
0
    lz->long_cnt = stepped_up;
3572
3573
0
    return result;
3574
0
}
3575
3576
static PyObject *
3577
count_next(PyObject *op)
3578
0
{
3579
0
    countobject *lz = countobject_CAST(op);
3580
0
#ifndef Py_GIL_DISABLED
3581
0
    if (lz->cnt == PY_SSIZE_T_MAX)
3582
0
        return count_nextlong(lz);
3583
0
    return PyLong_FromSsize_t(lz->cnt++);
3584
#else
3585
    // free-threading version
3586
    // fast mode uses compare-exchange loop
3587
    // slow mode uses a critical section
3588
    PyObject *returned;
3589
    Py_ssize_t cnt;
3590
3591
    cnt = _Py_atomic_load_ssize_relaxed(&lz->cnt);
3592
    for (;;) {
3593
        if (cnt == PY_SSIZE_T_MAX) {
3594
            Py_BEGIN_CRITICAL_SECTION(lz);
3595
            returned = count_nextlong(lz);
3596
            Py_END_CRITICAL_SECTION();
3597
            return returned;
3598
        }
3599
        if (_Py_atomic_compare_exchange_ssize(&lz->cnt, &cnt, cnt + 1)) {
3600
            return PyLong_FromSsize_t(cnt);
3601
        }
3602
    }
3603
#endif
3604
0
}
3605
3606
static PyObject *
3607
count_repr(PyObject *op)
3608
0
{
3609
0
    countobject *lz = countobject_CAST(op);
3610
0
    if (lz->long_cnt == NULL)
3611
0
        return PyUnicode_FromFormat("%s(%zd)",
3612
0
                                    _PyType_Name(Py_TYPE(lz)), lz->cnt);
3613
3614
0
    if (PyLong_Check(lz->long_step)) {
3615
0
        long step = PyLong_AsLong(lz->long_step);
3616
0
        if (step == -1 && PyErr_Occurred()) {
3617
0
            PyErr_Clear();
3618
0
        }
3619
0
        if (step == 1) {
3620
            /* Don't display step when it is an integer equal to 1 */
3621
0
            return PyUnicode_FromFormat("%s(%R)",
3622
0
                                        _PyType_Name(Py_TYPE(lz)),
3623
0
                                        lz->long_cnt);
3624
0
        }
3625
0
    }
3626
0
    return PyUnicode_FromFormat("%s(%R, %R)",
3627
0
                                _PyType_Name(Py_TYPE(lz)),
3628
0
                                lz->long_cnt, lz->long_step);
3629
0
}
3630
3631
static PyType_Slot count_slots[] = {
3632
    {Py_tp_dealloc, count_dealloc},
3633
    {Py_tp_repr, count_repr},
3634
    {Py_tp_getattro, PyObject_GenericGetAttr},
3635
    {Py_tp_doc, (void *)itertools_count__doc__},
3636
    {Py_tp_traverse, count_traverse},
3637
    {Py_tp_iter, PyObject_SelfIter},
3638
    {Py_tp_iternext, count_next},
3639
    {Py_tp_new, itertools_count},
3640
    {Py_tp_free, PyObject_GC_Del},
3641
    {0, NULL},
3642
};
3643
3644
static PyType_Spec count_spec = {
3645
    .name = "itertools.count",
3646
    .basicsize = sizeof(countobject),
3647
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3648
              Py_TPFLAGS_IMMUTABLETYPE),
3649
    .slots = count_slots,
3650
};
3651
3652
3653
/* repeat object ************************************************************/
3654
3655
typedef struct {
3656
    PyObject_HEAD
3657
    PyObject *element;
3658
    Py_ssize_t cnt;
3659
} repeatobject;
3660
3661
0
#define repeatobject_CAST(op)   ((repeatobject *)(op))
3662
3663
static PyObject *
3664
repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3665
0
{
3666
0
    repeatobject *ro;
3667
0
    PyObject *element;
3668
0
    Py_ssize_t cnt = -1, n_args;
3669
0
    static char *kwargs[] = {"object", "times", NULL};
3670
3671
0
    n_args = PyTuple_GET_SIZE(args);
3672
0
    if (kwds != NULL)
3673
0
        n_args += PyDict_GET_SIZE(kwds);
3674
0
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3675
0
                                     &element, &cnt))
3676
0
        return NULL;
3677
    /* Does user supply times argument? */
3678
0
    if (n_args == 2 && cnt < 0)
3679
0
        cnt = 0;
3680
3681
0
    ro = (repeatobject *)type->tp_alloc(type, 0);
3682
0
    if (ro == NULL)
3683
0
        return NULL;
3684
0
    ro->element = Py_NewRef(element);
3685
0
    ro->cnt = cnt;
3686
0
    return (PyObject *)ro;
3687
0
}
3688
3689
static void
3690
repeat_dealloc(PyObject *op)
3691
0
{
3692
0
    repeatobject *ro = repeatobject_CAST(op);
3693
0
    PyTypeObject *tp = Py_TYPE(ro);
3694
0
    PyObject_GC_UnTrack(ro);
3695
0
    Py_XDECREF(ro->element);
3696
0
    tp->tp_free(ro);
3697
0
    Py_DECREF(tp);
3698
0
}
3699
3700
static int
3701
repeat_traverse(PyObject *op, visitproc visit, void *arg)
3702
0
{
3703
0
    repeatobject *ro = repeatobject_CAST(op);
3704
0
    Py_VISIT(Py_TYPE(ro));
3705
0
    Py_VISIT(ro->element);
3706
0
    return 0;
3707
0
}
3708
3709
static PyObject *
3710
repeat_next(PyObject *op)
3711
0
{
3712
0
    repeatobject *ro = repeatobject_CAST(op);
3713
0
    Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(ro->cnt);
3714
0
    if (cnt == 0) {
3715
0
        return NULL;
3716
0
    }
3717
0
    if (cnt > 0) {
3718
0
        cnt--;
3719
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(ro->cnt, cnt);
3720
0
    }
3721
0
    return Py_NewRef(ro->element);
3722
0
}
3723
3724
static PyObject *
3725
repeat_repr(PyObject *op)
3726
0
{
3727
0
    repeatobject *ro = repeatobject_CAST(op);
3728
0
    if (ro->cnt == -1)
3729
0
        return PyUnicode_FromFormat("%s(%R)",
3730
0
                                    _PyType_Name(Py_TYPE(ro)), ro->element);
3731
0
    else
3732
0
        return PyUnicode_FromFormat("%s(%R, %zd)",
3733
0
                                    _PyType_Name(Py_TYPE(ro)), ro->element,
3734
0
                                    ro->cnt);
3735
0
}
3736
3737
static PyObject *
3738
repeat_len(PyObject *op, PyObject *Py_UNUSED(args))
3739
0
{
3740
0
    repeatobject *ro = repeatobject_CAST(op);
3741
0
    if (ro->cnt == -1) {
3742
0
        PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3743
0
        return NULL;
3744
0
    }
3745
0
    return PyLong_FromSize_t(ro->cnt);
3746
0
}
3747
3748
PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
3749
3750
static PyMethodDef repeat_methods[] = {
3751
    {"__length_hint__", repeat_len, METH_NOARGS, length_hint_doc},
3752
    {NULL,              NULL}           /* sentinel */
3753
};
3754
3755
PyDoc_STRVAR(repeat_doc,
3756
"repeat(object [,times]) -> create an iterator which returns the object\n\
3757
for the specified number of times.  If not specified, returns the object\n\
3758
endlessly.");
3759
3760
static PyType_Slot repeat_slots[] = {
3761
    {Py_tp_dealloc, repeat_dealloc},
3762
    {Py_tp_repr, repeat_repr},
3763
    {Py_tp_getattro, PyObject_GenericGetAttr},
3764
    {Py_tp_doc, (void *)repeat_doc},
3765
    {Py_tp_traverse, repeat_traverse},
3766
    {Py_tp_iter, PyObject_SelfIter},
3767
    {Py_tp_iternext, repeat_next},
3768
    {Py_tp_methods, repeat_methods},
3769
    {Py_tp_new, repeat_new},
3770
    {Py_tp_free, PyObject_GC_Del},
3771
    {0, NULL},
3772
};
3773
3774
static PyType_Spec repeat_spec = {
3775
    .name = "itertools.repeat",
3776
    .basicsize = sizeof(repeatobject),
3777
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3778
              Py_TPFLAGS_IMMUTABLETYPE),
3779
    .slots = repeat_slots,
3780
};
3781
3782
3783
/* ziplongest object *********************************************************/
3784
3785
typedef struct {
3786
    PyObject_HEAD
3787
    Py_ssize_t tuplesize;
3788
    Py_ssize_t numactive;
3789
    PyObject *ittuple;                  /* tuple of iterators */
3790
    PyObject *result;
3791
    PyObject *fillvalue;
3792
} ziplongestobject;
3793
3794
0
#define ziplongestobject_CAST(op)   ((ziplongestobject *)(op))
3795
3796
static PyObject *
3797
zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3798
0
{
3799
0
    ziplongestobject *lz;
3800
0
    Py_ssize_t i;
3801
0
    PyObject *ittuple;  /* tuple of iterators */
3802
0
    PyObject *result;
3803
0
    PyObject *fillvalue = Py_None;
3804
0
    Py_ssize_t tuplesize;
3805
3806
0
    if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
3807
0
        fillvalue = NULL;
3808
0
        if (PyDict_GET_SIZE(kwds) == 1) {
3809
0
            fillvalue = PyDict_GetItemWithError(kwds, &_Py_ID(fillvalue));
3810
0
        }
3811
0
        if (fillvalue == NULL) {
3812
0
            if (!PyErr_Occurred()) {
3813
0
                PyErr_SetString(PyExc_TypeError,
3814
0
                    "zip_longest() got an unexpected keyword argument");
3815
0
            }
3816
0
            return NULL;
3817
0
        }
3818
0
    }
3819
3820
    /* args must be a tuple */
3821
0
    assert(PyTuple_Check(args));
3822
0
    tuplesize = PyTuple_GET_SIZE(args);
3823
3824
    /* obtain iterators */
3825
0
    ittuple = PyTuple_New(tuplesize);
3826
0
    if (ittuple == NULL)
3827
0
        return NULL;
3828
0
    for (i=0; i < tuplesize; i++) {
3829
0
        PyObject *item = PyTuple_GET_ITEM(args, i);
3830
0
        PyObject *it = PyObject_GetIter(item);
3831
0
        if (it == NULL) {
3832
0
            Py_DECREF(ittuple);
3833
0
            return NULL;
3834
0
        }
3835
0
        PyTuple_SET_ITEM(ittuple, i, it);
3836
0
    }
3837
3838
    /* create a result holder */
3839
0
    result = PyTuple_New(tuplesize);
3840
0
    if (result == NULL) {
3841
0
        Py_DECREF(ittuple);
3842
0
        return NULL;
3843
0
    }
3844
0
    for (i=0 ; i < tuplesize ; i++) {
3845
0
        Py_INCREF(Py_None);
3846
0
        PyTuple_SET_ITEM(result, i, Py_None);
3847
0
    }
3848
3849
    /* create ziplongestobject structure */
3850
0
    lz = (ziplongestobject *)type->tp_alloc(type, 0);
3851
0
    if (lz == NULL) {
3852
0
        Py_DECREF(ittuple);
3853
0
        Py_DECREF(result);
3854
0
        return NULL;
3855
0
    }
3856
0
    lz->ittuple = ittuple;
3857
0
    lz->tuplesize = tuplesize;
3858
0
    lz->numactive = tuplesize;
3859
0
    lz->result = result;
3860
0
    lz->fillvalue = Py_NewRef(fillvalue);
3861
0
    return (PyObject *)lz;
3862
0
}
3863
3864
static void
3865
zip_longest_dealloc(PyObject *op)
3866
0
{
3867
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3868
0
    PyTypeObject *tp = Py_TYPE(lz);
3869
0
    PyObject_GC_UnTrack(lz);
3870
0
    Py_XDECREF(lz->ittuple);
3871
0
    Py_XDECREF(lz->result);
3872
0
    Py_XDECREF(lz->fillvalue);
3873
0
    tp->tp_free(lz);
3874
0
    Py_DECREF(tp);
3875
0
}
3876
3877
static int
3878
zip_longest_traverse(PyObject *op, visitproc visit, void *arg)
3879
0
{
3880
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3881
0
    Py_VISIT(Py_TYPE(lz));
3882
0
    Py_VISIT(lz->ittuple);
3883
0
    Py_VISIT(lz->result);
3884
0
    Py_VISIT(lz->fillvalue);
3885
0
    return 0;
3886
0
}
3887
3888
static PyObject *
3889
zip_longest_next_lock_held(PyObject *op)
3890
0
{
3891
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3892
0
    Py_ssize_t i;
3893
0
    Py_ssize_t tuplesize = lz->tuplesize;
3894
0
    PyObject *result = lz->result;
3895
0
    PyObject *it;
3896
0
    PyObject *item;
3897
0
    PyObject *olditem;
3898
3899
0
    if (tuplesize == 0)
3900
0
        return NULL;
3901
0
    if (lz->numactive == 0)
3902
0
        return NULL;
3903
0
    if (_PyObject_IsUniquelyReferenced(result)) {
3904
0
        Py_INCREF(result);
3905
0
        for (i=0 ; i < tuplesize ; i++) {
3906
0
            it = PyTuple_GET_ITEM(lz->ittuple, i);
3907
0
            if (it == NULL) {
3908
0
                item = Py_NewRef(lz->fillvalue);
3909
0
            } else {
3910
0
                item = PyIter_Next(it);
3911
0
                if (item == NULL) {
3912
0
                    lz->numactive -= 1;
3913
0
                    if (lz->numactive == 0 || PyErr_Occurred()) {
3914
0
                        lz->numactive = 0;
3915
0
                        Py_DECREF(result);
3916
0
                        return NULL;
3917
0
                    } else {
3918
0
                        item = Py_NewRef(lz->fillvalue);
3919
0
                        PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3920
0
                        Py_DECREF(it);
3921
0
                    }
3922
0
                }
3923
0
            }
3924
0
            olditem = PyTuple_GET_ITEM(result, i);
3925
0
            PyTuple_SET_ITEM(result, i, item);
3926
0
            Py_DECREF(olditem);
3927
0
        }
3928
        // bpo-42536: The GC may have untracked this result tuple. Since we're
3929
        // recycling it, make sure it's tracked again:
3930
0
        _PyTuple_Recycle(result);
3931
0
    } else {
3932
0
        result = PyTuple_New(tuplesize);
3933
0
        if (result == NULL)
3934
0
            return NULL;
3935
0
        for (i=0 ; i < tuplesize ; i++) {
3936
0
            it = PyTuple_GET_ITEM(lz->ittuple, i);
3937
0
            if (it == NULL) {
3938
0
                item = Py_NewRef(lz->fillvalue);
3939
0
            } else {
3940
0
                item = PyIter_Next(it);
3941
0
                if (item == NULL) {
3942
0
                    lz->numactive -= 1;
3943
0
                    if (lz->numactive == 0 || PyErr_Occurred()) {
3944
0
                        lz->numactive = 0;
3945
0
                        Py_DECREF(result);
3946
0
                        return NULL;
3947
0
                    } else {
3948
0
                        item = Py_NewRef(lz->fillvalue);
3949
0
                        PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3950
0
                        Py_DECREF(it);
3951
0
                    }
3952
0
                }
3953
0
            }
3954
0
            PyTuple_SET_ITEM(result, i, item);
3955
0
        }
3956
0
    }
3957
0
    return result;
3958
0
}
3959
3960
static PyObject *
3961
zip_longest_next(PyObject *op)
3962
0
{
3963
0
    PyObject *result;
3964
0
    Py_BEGIN_CRITICAL_SECTION(op);
3965
0
    result = zip_longest_next_lock_held(op);
3966
0
    Py_END_CRITICAL_SECTION()
3967
0
    return result;
3968
0
}
3969
3970
PyDoc_STRVAR(zip_longest_doc,
3971
"zip_longest(*iterables, fillvalue=None)\n\
3972
--\n\
3973
\n\
3974
Return a zip_longest object whose .__next__() method returns a tuple where\n\
3975
the i-th element comes from the i-th iterable argument.  The .__next__()\n\
3976
method continues until the longest iterable in the argument sequence\n\
3977
is exhausted and then it raises StopIteration.  When the shorter iterables\n\
3978
are exhausted, the fillvalue is substituted in their place.  The fillvalue\n\
3979
defaults to None or can be specified by a keyword argument.\n\
3980
");
3981
3982
static PyType_Slot ziplongest_slots[] = {
3983
    {Py_tp_dealloc, zip_longest_dealloc},
3984
    {Py_tp_getattro, PyObject_GenericGetAttr},
3985
    {Py_tp_doc, (void *)zip_longest_doc},
3986
    {Py_tp_traverse, zip_longest_traverse},
3987
    {Py_tp_iter, PyObject_SelfIter},
3988
    {Py_tp_iternext, zip_longest_next},
3989
    {Py_tp_new, zip_longest_new},
3990
    {Py_tp_free, PyObject_GC_Del},
3991
    {0, NULL},
3992
};
3993
3994
static PyType_Spec ziplongest_spec = {
3995
    .name = "itertools.zip_longest",
3996
    .basicsize = sizeof(ziplongestobject),
3997
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3998
              Py_TPFLAGS_IMMUTABLETYPE),
3999
    .slots = ziplongest_slots,
4000
};
4001
4002
4003
/* module level code ********************************************************/
4004
4005
PyDoc_STRVAR(module_doc,
4006
"Functional tools for creating and using iterators.\n\
4007
\n\
4008
Infinite iterators:\n\
4009
count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
4010
cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
4011
repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
4012
\n\
4013
Iterators terminating on the shortest input sequence:\n\
4014
accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
4015
batched(p, n) --> [p0, p1, ..., p_n-1], [p_n, p_n+1, ..., p_2n-1], ...\n\
4016
chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4017
chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
4018
compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4019
dropwhile(predicate, seq) --> seq[n], seq[n+1], starting when predicate fails\n\
4020
groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
4021
filterfalse(predicate, seq) --> elements of seq where predicate(elem) is False\n\
4022
islice(seq, [start,] stop [, step]) --> elements from\n\
4023
       seq[start:stop:step]\n\
4024
pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
4025
starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
4026
tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
4027
takewhile(predicate, seq) --> seq[0], seq[1], until predicate fails\n\
4028
zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
4029
\n\
4030
Combinatoric generators:\n\
4031
product(p, q, ... [repeat=1]) --> cartesian product\n\
4032
permutations(p[, r])\n\
4033
combinations(p, r)\n\
4034
combinations_with_replacement(p, r)\n\
4035
");
4036
4037
static int
4038
itertoolsmodule_traverse(PyObject *mod, visitproc visit, void *arg)
4039
60
{
4040
60
    itertools_state *state = get_module_state(mod);
4041
60
    Py_VISIT(state->accumulate_type);
4042
60
    Py_VISIT(state->batched_type);
4043
60
    Py_VISIT(state->chain_type);
4044
60
    Py_VISIT(state->combinations_type);
4045
60
    Py_VISIT(state->compress_type);
4046
60
    Py_VISIT(state->count_type);
4047
60
    Py_VISIT(state->cwr_type);
4048
60
    Py_VISIT(state->cycle_type);
4049
60
    Py_VISIT(state->dropwhile_type);
4050
60
    Py_VISIT(state->filterfalse_type);
4051
60
    Py_VISIT(state->groupby_type);
4052
60
    Py_VISIT(state->_grouper_type);
4053
60
    Py_VISIT(state->islice_type);
4054
60
    Py_VISIT(state->pairwise_type);
4055
60
    Py_VISIT(state->permutations_type);
4056
60
    Py_VISIT(state->product_type);
4057
60
    Py_VISIT(state->repeat_type);
4058
60
    Py_VISIT(state->starmap_type);
4059
60
    Py_VISIT(state->takewhile_type);
4060
60
    Py_VISIT(state->tee_type);
4061
60
    Py_VISIT(state->teedataobject_type);
4062
60
    Py_VISIT(state->ziplongest_type);
4063
60
    return 0;
4064
60
}
4065
4066
static int
4067
itertoolsmodule_clear(PyObject *mod)
4068
0
{
4069
0
    itertools_state *state = get_module_state(mod);
4070
0
    Py_CLEAR(state->accumulate_type);
4071
0
    Py_CLEAR(state->batched_type);
4072
0
    Py_CLEAR(state->chain_type);
4073
0
    Py_CLEAR(state->combinations_type);
4074
0
    Py_CLEAR(state->compress_type);
4075
0
    Py_CLEAR(state->count_type);
4076
0
    Py_CLEAR(state->cwr_type);
4077
0
    Py_CLEAR(state->cycle_type);
4078
0
    Py_CLEAR(state->dropwhile_type);
4079
0
    Py_CLEAR(state->filterfalse_type);
4080
0
    Py_CLEAR(state->groupby_type);
4081
0
    Py_CLEAR(state->_grouper_type);
4082
0
    Py_CLEAR(state->islice_type);
4083
0
    Py_CLEAR(state->pairwise_type);
4084
0
    Py_CLEAR(state->permutations_type);
4085
0
    Py_CLEAR(state->product_type);
4086
0
    Py_CLEAR(state->repeat_type);
4087
0
    Py_CLEAR(state->starmap_type);
4088
0
    Py_CLEAR(state->takewhile_type);
4089
0
    Py_CLEAR(state->tee_type);
4090
0
    Py_CLEAR(state->teedataobject_type);
4091
0
    Py_CLEAR(state->ziplongest_type);
4092
0
    return 0;
4093
0
}
4094
4095
static void
4096
itertoolsmodule_free(void *mod)
4097
0
{
4098
0
    (void)itertoolsmodule_clear((PyObject *)mod);
4099
0
}
4100
4101
110
#define ADD_TYPE(module, type, spec)                                     \
4102
110
do {                                                                     \
4103
110
    type = (PyTypeObject *)PyType_FromModuleAndSpec(module, spec, NULL); \
4104
110
    if (type == NULL) {                                                  \
4105
0
        return -1;                                                       \
4106
0
    }                                                                    \
4107
110
    if (PyModule_AddType(module, type) < 0) {                            \
4108
0
        return -1;                                                       \
4109
0
    }                                                                    \
4110
110
} while (0)
4111
4112
static int
4113
itertoolsmodule_exec(PyObject *mod)
4114
5
{
4115
5
    itertools_state *state = get_module_state(mod);
4116
5
    ADD_TYPE(mod, state->accumulate_type, &accumulate_spec);
4117
5
    ADD_TYPE(mod, state->batched_type, &batched_spec);
4118
5
    ADD_TYPE(mod, state->chain_type, &chain_spec);
4119
5
    ADD_TYPE(mod, state->combinations_type, &combinations_spec);
4120
5
    ADD_TYPE(mod, state->compress_type, &compress_spec);
4121
5
    ADD_TYPE(mod, state->count_type, &count_spec);
4122
5
    ADD_TYPE(mod, state->cwr_type, &cwr_spec);
4123
5
    ADD_TYPE(mod, state->cycle_type, &cycle_spec);
4124
5
    ADD_TYPE(mod, state->dropwhile_type, &dropwhile_spec);
4125
5
    ADD_TYPE(mod, state->filterfalse_type, &filterfalse_spec);
4126
5
    ADD_TYPE(mod, state->groupby_type, &groupby_spec);
4127
5
    ADD_TYPE(mod, state->_grouper_type, &_grouper_spec);
4128
5
    ADD_TYPE(mod, state->islice_type, &islice_spec);
4129
5
    ADD_TYPE(mod, state->pairwise_type, &pairwise_spec);
4130
5
    ADD_TYPE(mod, state->permutations_type, &permutations_spec);
4131
5
    ADD_TYPE(mod, state->product_type, &product_spec);
4132
5
    ADD_TYPE(mod, state->repeat_type, &repeat_spec);
4133
5
    ADD_TYPE(mod, state->starmap_type, &starmap_spec);
4134
5
    ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
4135
5
    ADD_TYPE(mod, state->tee_type, &tee_spec);
4136
5
    ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
4137
5
    ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
4138
4139
5
    Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
4140
5
    return 0;
4141
5
}
4142
4143
static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4144
    _Py_ABI_SLOT,
4145
    {Py_mod_exec, itertoolsmodule_exec},
4146
    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
4147
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
4148
    {0, NULL}
4149
};
4150
4151
static PyMethodDef module_methods[] = {
4152
    ITERTOOLS_TEE_METHODDEF
4153
    {NULL, NULL} /* sentinel */
4154
};
4155
4156
4157
static struct PyModuleDef itertoolsmodule = {
4158
    .m_base = PyModuleDef_HEAD_INIT,
4159
    .m_name = "itertools",
4160
    .m_doc = module_doc,
4161
    .m_size = sizeof(itertools_state),
4162
    .m_methods = module_methods,
4163
    .m_slots = itertoolsmodule_slots,
4164
    .m_traverse = itertoolsmodule_traverse,
4165
    .m_clear = itertoolsmodule_clear,
4166
    .m_free = itertoolsmodule_free,
4167
};
4168
4169
PyMODINIT_FUNC
4170
PyInit_itertools(void)
4171
5
{
4172
5
    return PyModuleDef_Init(&itertoolsmodule);
4173
5
}