Coverage Report

Created: 2026-07-14 06:16

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