Coverage Report

Created: 2026-08-13 06:33

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/src/cpython3/Modules/itertoolsmodule.c
Line
Count
Source
1
#include "Python.h"
2
#include "pycore_call.h"              // _PyObject_CallNoArgs()
3
#include "pycore_ceval.h"             // _PyEval_GetBuiltin()
4
#include "pycore_critical_section.h"  // Py_BEGIN_CRITICAL_SECTION()
5
#include "pycore_long.h"              // _PyLong_GetZero()
6
#include "pycore_moduleobject.h"      // _PyModule_GetState()
7
#include "pycore_typeobject.h"        // _PyType_GetModuleState()
8
#include "pycore_object.h"            // _PyObject_GC_TRACK()
9
#include "pycore_tuple.h"             // _PyTuple_ITEMS()
10
11
#include <stddef.h>                   // offsetof()
12
13
/* Itertools module written and maintained
14
   by Raymond D. Hettinger <python@rcn.com>
15
*/
16
17
typedef struct {
18
    PyTypeObject *accumulate_type;
19
    PyTypeObject *batched_type;
20
    PyTypeObject *chain_type;
21
    PyTypeObject *combinations_type;
22
    PyTypeObject *compress_type;
23
    PyTypeObject *count_type;
24
    PyTypeObject *cwr_type;
25
    PyTypeObject *cycle_type;
26
    PyTypeObject *dropwhile_type;
27
    PyTypeObject *filterfalse_type;
28
    PyTypeObject *groupby_type;
29
    PyTypeObject *_grouper_type;
30
    PyTypeObject *islice_type;
31
    PyTypeObject *pairwise_type;
32
    PyTypeObject *permutations_type;
33
    PyTypeObject *product_type;
34
    PyTypeObject *repeat_type;
35
    PyTypeObject *starmap_type;
36
    PyTypeObject *takewhile_type;
37
    PyTypeObject *tee_type;
38
    PyTypeObject *teedataobject_type;
39
    PyTypeObject *ziplongest_type;
40
} itertools_state;
41
42
static inline itertools_state *
43
get_module_state(PyObject *mod)
44
362
{
45
362
    void *state = _PyModule_GetState(mod);
46
362
    assert(state != NULL);
47
362
    return (itertools_state *)state;
48
362
}
49
50
static inline itertools_state *
51
get_module_state_by_cls(PyTypeObject *cls)
52
0
{
53
0
    void *state = _PyType_GetModuleState(cls);
54
0
    assert(state != NULL);
55
0
    return (itertools_state *)state;
56
0
}
57
58
static struct PyModuleDef itertoolsmodule;
59
60
static inline itertools_state *
61
find_state_by_type(PyTypeObject *tp)
62
0
{
63
0
    PyObject *mod = PyType_GetModuleByDef(tp, &itertoolsmodule);
64
0
    assert(mod != NULL);
65
0
    return get_module_state(mod);
66
0
}
67
68
/*[clinic input]
69
module itertools
70
class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
71
class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
72
class itertools.teedataobject "teedataobject *" "clinic_state()->teedataobject_type"
73
class itertools._tee "teeobject *" "clinic_state()->tee_type"
74
class itertools.batched "batchedobject *" "clinic_state()->batched_type"
75
class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
76
class itertools.dropwhile "dropwhileobject *" "clinic_state()->dropwhile_type"
77
class itertools.takewhile "takewhileobject *" "clinic_state()->takewhile_type"
78
class itertools.starmap "starmapobject *" "clinic_state()->starmap_type"
79
class itertools.chain "chainobject *" "clinic_state()->chain_type"
80
class itertools.combinations "combinationsobject *" "clinic_state()->combinations_type"
81
class itertools.combinations_with_replacement "cwr_object *" "clinic_state()->cwr_type"
82
class itertools.permutations "permutationsobject *" "clinic_state()->permutations_type"
83
class itertools.accumulate "accumulateobject *" "clinic_state()->accumulate_type"
84
class itertools.compress "compressobject *" "clinic_state()->compress_type"
85
class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_type"
86
class itertools.count "countobject *" "clinic_state()->count_type"
87
class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
88
[clinic start generated code]*/
89
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=aa48fe4de9d4080f]*/
90
91
0
#define clinic_state() (find_state_by_type(type))
92
0
#define clinic_state_by_cls() (get_module_state_by_cls(base_tp))
93
#include "clinic/itertoolsmodule.c.h"
94
#undef clinic_state_by_cls
95
#undef clinic_state
96
97
98
/* batched object ************************************************************/
99
100
typedef struct {
101
    PyObject_HEAD
102
    PyObject *it;
103
    Py_ssize_t batch_size;
104
    bool strict;
105
} batchedobject;
106
107
0
#define batchedobject_CAST(op)  ((batchedobject *)(op))
108
109
/*[clinic input]
110
@permit_long_summary
111
@classmethod
112
itertools.batched.__new__ as batched_new
113
    iterable: object
114
    n: Py_ssize_t
115
    *
116
    strict: bool = False
117
118
Batch data into tuples of length n. The last batch may be shorter than n.
119
120
Loops over the input iterable and accumulates data into tuples
121
up to size n.  The input is consumed lazily, just enough to
122
fill a batch.  The result is yielded as soon as a batch is full
123
or when the input iterable is exhausted.
124
125
    >>> for batch in batched('ABCDEFG', 3):
126
    ...     print(batch)
127
    ...
128
    ('A', 'B', 'C')
129
    ('D', 'E', 'F')
130
    ('G',)
131
132
If "strict" is True, raises a ValueError if the final batch is shorter
133
than n.
134
135
[clinic start generated code]*/
136
137
static PyObject *
138
batched_new_impl(PyTypeObject *type, PyObject *iterable, Py_ssize_t n,
139
                 int strict)
140
/*[clinic end generated code: output=c6de11b061529d3e input=b31d8be8e8577a34]*/
141
0
{
142
0
    PyObject *it;
143
0
    batchedobject *bo;
144
145
0
    if (n < 1) {
146
        /* We could define the n==0 case to return an empty iterator
147
           but that is at odds with the idea that batching should
148
           never throw-away input data.
149
        */
150
0
        PyErr_SetString(PyExc_ValueError, "n must be at least one");
151
0
        return NULL;
152
0
    }
153
0
    it = PyObject_GetIter(iterable);
154
0
    if (it == NULL) {
155
0
        return NULL;
156
0
    }
157
158
    /* create batchedobject structure */
159
0
    bo = (batchedobject *)type->tp_alloc(type, 0);
160
0
    if (bo == NULL) {
161
0
        Py_DECREF(it);
162
0
        return NULL;
163
0
    }
164
0
    bo->batch_size = n;
165
0
    bo->it = it;
166
0
    bo->strict = (bool) strict;
167
0
    return (PyObject *)bo;
168
0
}
169
170
static void
171
batched_dealloc(PyObject *op)
172
0
{
173
0
    batchedobject *bo = batchedobject_CAST(op);
174
0
    PyTypeObject *tp = Py_TYPE(bo);
175
0
    PyObject_GC_UnTrack(bo);
176
0
    Py_XDECREF(bo->it);
177
0
    tp->tp_free(bo);
178
0
    Py_DECREF(tp);
179
0
}
180
181
static int
182
batched_traverse(PyObject *op, visitproc visit, void *arg)
183
0
{
184
0
    batchedobject *bo = batchedobject_CAST(op);
185
0
    Py_VISIT(Py_TYPE(bo));
186
0
    Py_VISIT(bo->it);
187
0
    return 0;
188
0
}
189
190
static PyObject *
191
batched_next(PyObject *op)
192
0
{
193
0
    batchedobject *bo = batchedobject_CAST(op);
194
0
    Py_ssize_t i;
195
0
    Py_ssize_t n = FT_ATOMIC_LOAD_SSIZE_RELAXED(bo->batch_size);
196
0
    PyObject *it = bo->it;
197
0
    PyObject *item;
198
0
    PyObject *result;
199
200
0
    if (n < 0) {
201
0
        return NULL;
202
0
    }
203
0
    result = PyTuple_New(n);
204
0
    if (result == NULL) {
205
0
        return NULL;
206
0
    }
207
0
    iternextfunc iternext = *Py_TYPE(it)->tp_iternext;
208
0
    PyObject **items = _PyTuple_ITEMS(result);
209
0
    for (i=0 ; i < n ; i++) {
210
0
        item = iternext(it);
211
0
        if (item == NULL) {
212
0
            goto null_item;
213
0
        }
214
0
        items[i] = item;
215
0
    }
216
0
    return result;
217
218
0
 null_item:
219
0
    if (PyErr_Occurred()) {
220
0
        if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
221
            /* Input raised an exception other than StopIteration */
222
0
            FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
223
0
#ifndef Py_GIL_DISABLED
224
0
            Py_CLEAR(bo->it);
225
0
#endif
226
0
            Py_DECREF(result);
227
0
            return NULL;
228
0
        }
229
0
        PyErr_Clear();
230
0
    }
231
0
    if (i == 0) {
232
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
233
0
#ifndef Py_GIL_DISABLED
234
0
        Py_CLEAR(bo->it);
235
0
#endif
236
0
        Py_DECREF(result);
237
0
        return NULL;
238
0
    }
239
0
    if (bo->strict) {
240
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
241
0
#ifndef Py_GIL_DISABLED
242
0
        Py_CLEAR(bo->it);
243
0
#endif
244
0
        Py_DECREF(result);
245
0
        PyErr_SetString(PyExc_ValueError, "batched(): incomplete batch");
246
0
        return NULL;
247
0
    }
248
0
    _PyTuple_Resize(&result, i);
249
0
    return result;
250
0
}
251
252
static PyType_Slot batched_slots[] = {
253
    {Py_tp_dealloc, batched_dealloc},
254
    {Py_tp_getattro, PyObject_GenericGetAttr},
255
    {Py_tp_doc, (void *)batched_new__doc__},
256
    {Py_tp_traverse, batched_traverse},
257
    {Py_tp_iter, PyObject_SelfIter},
258
    {Py_tp_iternext, batched_next},
259
    {Py_tp_alloc, PyType_GenericAlloc},
260
    {Py_tp_new, batched_new},
261
    {Py_tp_free, PyObject_GC_Del},
262
    {0, NULL},
263
};
264
265
static PyType_Spec batched_spec = {
266
    .name = "itertools.batched",
267
    .basicsize = sizeof(batchedobject),
268
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
269
              Py_TPFLAGS_IMMUTABLETYPE),
270
    .slots = batched_slots,
271
};
272
273
274
/* pairwise object ***********************************************************/
275
276
typedef struct {
277
    PyObject_HEAD
278
    PyObject *it;
279
    PyObject *old;
280
    PyObject *result;
281
} pairwiseobject;
282
283
0
#define pairwiseobject_CAST(op) ((pairwiseobject *)(op))
284
285
/*[clinic input]
286
@classmethod
287
itertools.pairwise.__new__ as pairwise_new
288
    iterable: object
289
    /
290
Return an iterator of overlapping pairs taken from the input iterator.
291
292
    s -> (s0,s1), (s1,s2), (s2, s3), ...
293
294
[clinic start generated code]*/
295
296
static PyObject *
297
pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
298
/*[clinic end generated code: output=9f0267062d384456 input=6e7c3cddb431a8d6]*/
299
0
{
300
0
    PyObject *it;
301
0
    pairwiseobject *po;
302
303
0
    it = PyObject_GetIter(iterable);
304
0
    if (it == NULL) {
305
0
        return NULL;
306
0
    }
307
0
    po = (pairwiseobject *)type->tp_alloc(type, 0);
308
0
    if (po == NULL) {
309
0
        Py_DECREF(it);
310
0
        return NULL;
311
0
    }
312
0
    po->it = it;
313
0
    po->old = NULL;
314
0
    po->result = _PyTuple_FromPairSteal(Py_None, Py_None);
315
0
    if (po->result == NULL) {
316
0
        Py_DECREF(po);
317
0
        return NULL;
318
0
    }
319
0
    return (PyObject *)po;
320
0
}
321
322
static void
323
pairwise_dealloc(PyObject *op)
324
0
{
325
0
    pairwiseobject *po = pairwiseobject_CAST(op);
326
0
    PyTypeObject *tp = Py_TYPE(po);
327
0
    PyObject_GC_UnTrack(po);
328
0
    Py_XDECREF(po->it);
329
0
    Py_XDECREF(po->old);
330
0
    Py_XDECREF(po->result);
331
0
    tp->tp_free(po);
332
0
    Py_DECREF(tp);
333
0
}
334
335
static int
336
pairwise_traverse(PyObject *op, visitproc visit, void *arg)
337
0
{
338
0
    pairwiseobject *po = pairwiseobject_CAST(op);
339
0
    Py_VISIT(Py_TYPE(po));
340
0
    Py_VISIT(po->it);
341
0
    Py_VISIT(po->old);
342
0
    Py_VISIT(po->result);
343
0
    return 0;
344
0
}
345
346
static PyObject *
347
pairwise_next(PyObject *op)
348
0
{
349
0
    pairwiseobject *po = pairwiseobject_CAST(op);
350
0
    PyObject *it = po->it;
351
0
    PyObject *old = po->old;
352
0
    PyObject *new, *result;
353
354
0
    if (it == NULL) {
355
0
        return NULL;
356
0
    }
357
0
    if (old == NULL) {
358
0
        old = (*Py_TYPE(it)->tp_iternext)(it);
359
0
        Py_XSETREF(po->old, old);
360
0
        if (old == NULL) {
361
0
            Py_CLEAR(po->it);
362
0
            return NULL;
363
0
        }
364
0
        it = po->it;
365
0
        if (it == NULL) {
366
0
            Py_CLEAR(po->old);
367
0
            return NULL;
368
0
        }
369
0
    }
370
0
    Py_INCREF(old);
371
0
    new = (*Py_TYPE(it)->tp_iternext)(it);
372
0
    if (new == NULL) {
373
0
        Py_CLEAR(po->it);
374
0
        Py_CLEAR(po->old);
375
0
        Py_DECREF(old);
376
0
        return NULL;
377
0
    }
378
379
0
    result = po->result;
380
0
    if (_PyObject_IsUniquelyReferenced(result)) {
381
0
        Py_INCREF(result);
382
0
        PyObject *last_old = PyTuple_GET_ITEM(result, 0);
383
0
        PyObject *last_new = PyTuple_GET_ITEM(result, 1);
384
0
        PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
385
0
        PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
386
0
        Py_DECREF(last_old);
387
0
        Py_DECREF(last_new);
388
        // bpo-42536: The GC may have untracked this result tuple. Since we're
389
        // recycling it, make sure it's tracked again:
390
0
        _PyTuple_Recycle(result);
391
0
    }
392
0
    else {
393
0
        result = _PyTuple_FromPair(old, new);
394
0
    }
395
396
0
    Py_XSETREF(po->old, new);
397
0
    Py_DECREF(old);
398
0
    return result;
399
0
}
400
401
static PyType_Slot pairwise_slots[] = {
402
    {Py_tp_dealloc, pairwise_dealloc},
403
    {Py_tp_getattro, PyObject_GenericGetAttr},
404
    {Py_tp_doc, (void *)pairwise_new__doc__},
405
    {Py_tp_traverse, pairwise_traverse},
406
    {Py_tp_iter, PyObject_SelfIter},
407
    {Py_tp_iternext, pairwise_next},
408
    {Py_tp_alloc, PyType_GenericAlloc},
409
    {Py_tp_new, pairwise_new},
410
    {Py_tp_free, PyObject_GC_Del},
411
    {0, NULL},
412
};
413
414
static PyType_Spec pairwise_spec = {
415
    .name = "itertools.pairwise",
416
    .basicsize = sizeof(pairwiseobject),
417
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
418
              Py_TPFLAGS_IMMUTABLETYPE),
419
    .slots = pairwise_slots,
420
};
421
422
423
/* groupby object ************************************************************/
424
425
typedef struct {
426
    PyObject_HEAD
427
    PyObject *it;
428
    PyObject *keyfunc;
429
    PyObject *tgtkey;
430
    PyObject *currkey;
431
    PyObject *currvalue;
432
    const void *currgrouper;  /* borrowed reference */
433
    itertools_state *state;
434
} groupbyobject;
435
436
0
#define groupbyobject_CAST(op)  ((groupbyobject *)(op))
437
438
static PyObject *_grouper_create(groupbyobject *, PyObject *);
439
440
/*[clinic input]
441
@permit_long_summary
442
@classmethod
443
itertools.groupby.__new__
444
445
    iterable as it: object
446
        Elements to divide into groups according to the key function.
447
    key as keyfunc: object = None
448
        A function for computing the group category for each element.
449
        If the key function is not specified or is None, the element itself
450
        is used for grouping.
451
452
make an iterator that returns consecutive keys and groups from the iterable
453
[clinic start generated code]*/
454
455
static PyObject *
456
itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
457
/*[clinic end generated code: output=cbb1ae3a90fd4141 input=9f89fe625b20ef1a]*/
458
0
{
459
0
    groupbyobject *gbo;
460
461
0
    gbo = (groupbyobject *)type->tp_alloc(type, 0);
462
0
    if (gbo == NULL)
463
0
        return NULL;
464
0
    gbo->tgtkey = NULL;
465
0
    gbo->currkey = NULL;
466
0
    gbo->currvalue = NULL;
467
0
    gbo->keyfunc = Py_NewRef(keyfunc);
468
0
    gbo->it = PyObject_GetIter(it);
469
0
    if (gbo->it == NULL) {
470
0
        Py_DECREF(gbo);
471
0
        return NULL;
472
0
    }
473
0
    gbo->state = find_state_by_type(type);
474
0
    return (PyObject *)gbo;
475
0
}
476
477
static void
478
groupby_dealloc(PyObject *op)
479
0
{
480
0
    groupbyobject *gbo = groupbyobject_CAST(op);
481
0
    PyTypeObject *tp = Py_TYPE(gbo);
482
0
    PyObject_GC_UnTrack(gbo);
483
0
    Py_XDECREF(gbo->it);
484
0
    Py_XDECREF(gbo->keyfunc);
485
0
    Py_XDECREF(gbo->tgtkey);
486
0
    Py_XDECREF(gbo->currkey);
487
0
    Py_XDECREF(gbo->currvalue);
488
0
    tp->tp_free(gbo);
489
0
    Py_DECREF(tp);
490
0
}
491
492
static int
493
groupby_traverse(PyObject *op, visitproc visit, void *arg)
494
0
{
495
0
    groupbyobject *gbo = groupbyobject_CAST(op);
496
0
    Py_VISIT(Py_TYPE(gbo));
497
0
    Py_VISIT(gbo->it);
498
0
    Py_VISIT(gbo->keyfunc);
499
0
    Py_VISIT(gbo->tgtkey);
500
0
    Py_VISIT(gbo->currkey);
501
0
    Py_VISIT(gbo->currvalue);
502
0
    return 0;
503
0
}
504
505
Py_LOCAL_INLINE(int)
506
groupby_step(groupbyobject *gbo)
507
0
{
508
0
    PyObject *newvalue, *newkey, *oldvalue;
509
510
0
    newvalue = PyIter_Next(gbo->it);
511
0
    if (newvalue == NULL)
512
0
        return -1;
513
514
0
    if (gbo->keyfunc == Py_None) {
515
0
        newkey = Py_NewRef(newvalue);
516
0
    } else {
517
0
        newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
518
0
        if (newkey == NULL) {
519
0
            Py_DECREF(newvalue);
520
0
            return -1;
521
0
        }
522
0
    }
523
524
0
    oldvalue = gbo->currvalue;
525
0
    gbo->currvalue = newvalue;
526
0
    Py_XSETREF(gbo->currkey, newkey);
527
0
    Py_XDECREF(oldvalue);
528
0
    return 0;
529
0
}
530
531
static PyObject *
532
groupby_next(PyObject *op)
533
0
{
534
0
    PyObject *grouper;
535
0
    groupbyobject *gbo = groupbyobject_CAST(op);
536
537
0
    gbo->currgrouper = NULL;
538
    /* skip to next iteration group */
539
0
    for (;;) {
540
0
        if (gbo->currkey == NULL)
541
0
            /* pass */;
542
0
        else if (gbo->tgtkey == NULL)
543
0
            break;
544
0
        else {
545
            /* A user-defined __eq__ can re-enter groupby and advance the iterator,
546
               mutating gbo->tgtkey / gbo->currkey while we are comparing them.
547
               Take local snapshots and hold strong references so INCREF/DECREF
548
               apply to the same objects even under re-entrancy. */
549
0
            PyObject *tgtkey = gbo->tgtkey;
550
0
            PyObject *currkey = gbo->currkey;
551
552
0
            Py_INCREF(tgtkey);
553
0
            Py_INCREF(currkey);
554
0
            int rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
555
0
            Py_DECREF(tgtkey);
556
0
            Py_DECREF(currkey);
557
558
0
            if (rcmp == -1)
559
0
                return NULL;
560
0
            else if (rcmp == 0)
561
0
                break;
562
0
        }
563
564
0
        if (groupby_step(gbo) < 0)
565
0
            return NULL;
566
0
    }
567
0
    Py_INCREF(gbo->currkey);
568
0
    Py_XSETREF(gbo->tgtkey, gbo->currkey);
569
570
0
    grouper = _grouper_create(gbo, gbo->tgtkey);
571
0
    if (grouper == NULL)
572
0
        return NULL;
573
574
0
    return _PyTuple_FromPairSteal(Py_NewRef(gbo->currkey), grouper);
575
0
}
576
577
static PyType_Slot groupby_slots[] = {
578
    {Py_tp_dealloc, groupby_dealloc},
579
    {Py_tp_getattro, PyObject_GenericGetAttr},
580
    {Py_tp_doc, (void *)itertools_groupby__doc__},
581
    {Py_tp_traverse, groupby_traverse},
582
    {Py_tp_iter, PyObject_SelfIter},
583
    {Py_tp_iternext, groupby_next},
584
    {Py_tp_new, itertools_groupby},
585
    {Py_tp_free, PyObject_GC_Del},
586
    {0, NULL},
587
};
588
589
static PyType_Spec groupby_spec = {
590
    .name = "itertools.groupby",
591
    .basicsize= sizeof(groupbyobject),
592
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
593
              Py_TPFLAGS_IMMUTABLETYPE),
594
    .slots = groupby_slots,
595
};
596
597
/* _grouper object (internal) ************************************************/
598
599
typedef struct {
600
    PyObject_HEAD
601
    PyObject *parent;
602
    PyObject *tgtkey;
603
} _grouperobject;
604
605
0
#define _grouperobject_CAST(op) ((_grouperobject *)(op))
606
607
/*[clinic input]
608
@classmethod
609
itertools._grouper.__new__
610
611
    parent: object(subclass_of='clinic_state_by_cls()->groupby_type')
612
    tgtkey: object
613
    /
614
[clinic start generated code]*/
615
616
static PyObject *
617
itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
618
                        PyObject *tgtkey)
619
/*[clinic end generated code: output=462efb1cdebb5914 input=afe05eb477118f12]*/
620
0
{
621
0
    return _grouper_create(groupbyobject_CAST(parent), tgtkey);
622
0
}
623
624
static PyObject *
625
_grouper_create(groupbyobject *parent, PyObject *tgtkey)
626
0
{
627
0
    itertools_state *state = parent->state;
628
0
    _grouperobject *igo = PyObject_GC_New(_grouperobject, state->_grouper_type);
629
0
    if (igo == NULL)
630
0
        return NULL;
631
0
    igo->parent = Py_NewRef(parent);
632
0
    igo->tgtkey = Py_NewRef(tgtkey);
633
0
    parent->currgrouper = igo;  /* borrowed reference */
634
635
0
    PyObject_GC_Track(igo);
636
0
    return (PyObject *)igo;
637
0
}
638
639
static void
640
_grouper_dealloc(PyObject *op)
641
0
{
642
0
    _grouperobject *igo = _grouperobject_CAST(op);
643
0
    PyTypeObject *tp = Py_TYPE(igo);
644
0
    PyObject_GC_UnTrack(igo);
645
0
    Py_DECREF(igo->parent);
646
0
    Py_DECREF(igo->tgtkey);
647
0
    PyObject_GC_Del(igo);
648
0
    Py_DECREF(tp);
649
0
}
650
651
static int
652
_grouper_traverse(PyObject *op, visitproc visit, void *arg)
653
0
{
654
0
    _grouperobject *igo = _grouperobject_CAST(op);
655
0
    Py_VISIT(Py_TYPE(igo));
656
0
    Py_VISIT(igo->parent);
657
0
    Py_VISIT(igo->tgtkey);
658
0
    return 0;
659
0
}
660
661
static PyObject *
662
_grouper_next(PyObject *op)
663
0
{
664
0
    _grouperobject *igo = _grouperobject_CAST(op);
665
0
    groupbyobject *gbo = groupbyobject_CAST(igo->parent);
666
0
    PyObject *r;
667
0
    int rcmp;
668
669
0
    if (gbo->currgrouper != igo)
670
0
        return NULL;
671
0
    if (gbo->currvalue == NULL) {
672
0
        if (groupby_step(gbo) < 0)
673
0
            return NULL;
674
0
    }
675
676
0
    assert(gbo->currkey != NULL);
677
    /* A user-defined __eq__ can re-enter the grouper and advance the iterator,
678
       mutating gbo->currkey while we are comparing them.
679
       Take local snapshots and hold strong references so INCREF/DECREF
680
       apply to the same objects even under re-entrancy. */
681
0
    PyObject *tgtkey = Py_NewRef(igo->tgtkey);
682
0
    PyObject *currkey = Py_NewRef(gbo->currkey);
683
0
    rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
684
0
    Py_DECREF(tgtkey);
685
0
    Py_DECREF(currkey);
686
687
0
    if (rcmp <= 0)
688
        /* got any error or current group is end */
689
0
        return NULL;
690
691
0
    r = gbo->currvalue;
692
0
    gbo->currvalue = NULL;
693
0
    Py_CLEAR(gbo->currkey);
694
695
0
    return r;
696
0
}
697
698
static PyType_Slot _grouper_slots[] = {
699
    {Py_tp_dealloc, _grouper_dealloc},
700
    {Py_tp_getattro, PyObject_GenericGetAttr},
701
    {Py_tp_traverse, _grouper_traverse},
702
    {Py_tp_iter, PyObject_SelfIter},
703
    {Py_tp_iternext, _grouper_next},
704
    {Py_tp_new, itertools__grouper},
705
    {Py_tp_free, PyObject_GC_Del},
706
    {0, NULL},
707
};
708
709
static PyType_Spec _grouper_spec = {
710
    .name = "itertools._grouper",
711
    .basicsize = sizeof(_grouperobject),
712
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
713
              Py_TPFLAGS_IMMUTABLETYPE),
714
    .slots = _grouper_slots,
715
};
716
717
718
/* tee object and with supporting function and objects ***********************/
719
720
/* The teedataobject pre-allocates space for LINKCELLS number of objects.
721
   To help the object fit neatly inside cache lines (space for 16 to 32
722
   pointers), the value should be a multiple of 16 minus  space for
723
   the other structure members including PyHEAD overhead.  The larger the
724
   value, the less memory overhead per object and the less time spent
725
   allocating/deallocating new links.  The smaller the number, the less
726
   wasted space and the more rapid freeing of older data.
727
*/
728
0
#define LINKCELLS 57
729
730
typedef struct {
731
    PyObject_HEAD
732
    PyObject *it;
733
    int numread;                /* 0 <= numread <= LINKCELLS */
734
    int running;
735
    PyObject *nextlink;
736
    PyObject *(values[LINKCELLS]);
737
} teedataobject;
738
739
0
#define teedataobject_CAST(op)  ((teedataobject *)(op))
740
741
typedef struct {
742
    PyObject_HEAD
743
    teedataobject *dataobj;
744
    int index;                  /* 0 <= index <= LINKCELLS */
745
    PyObject *weakreflist;
746
    itertools_state *state;
747
} teeobject;
748
749
0
#define teeobject_CAST(op)  ((teeobject *)(op))
750
751
static PyObject *
752
teedataobject_newinternal(itertools_state *state, PyObject *it)
753
0
{
754
0
    teedataobject *tdo;
755
756
0
    tdo = PyObject_GC_New(teedataobject, state->teedataobject_type);
757
0
    if (tdo == NULL)
758
0
        return NULL;
759
760
0
    tdo->running = 0;
761
0
    tdo->numread = 0;
762
0
    tdo->nextlink = NULL;
763
0
    tdo->it = Py_NewRef(it);
764
0
    PyObject_GC_Track(tdo);
765
0
    return (PyObject *)tdo;
766
0
}
767
768
static PyObject *
769
teedataobject_jumplink(itertools_state *state, teedataobject *tdo)
770
0
{
771
0
    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
0
#define chainobject_CAST(op)    ((chainobject *)(op))
1887
1888
static PyObject *
1889
chain_new_internal(PyTypeObject *type, PyObject *source)
1890
0
{
1891
0
    chainobject *lz;
1892
1893
0
    lz = (chainobject *)type->tp_alloc(type, 0);
1894
0
    if (lz == NULL) {
1895
0
        Py_DECREF(source);
1896
0
        return NULL;
1897
0
    }
1898
1899
0
    lz->source = source;
1900
0
    lz->active = NULL;
1901
0
    return (PyObject *)lz;
1902
0
}
1903
1904
static PyObject *
1905
chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1906
0
{
1907
0
    PyObject *source;
1908
1909
0
    itertools_state *state = find_state_by_type(type);
1910
0
    PyTypeObject *chain_type = state->chain_type;
1911
0
    if ((type == chain_type || type->tp_init == chain_type->tp_init) &&
1912
0
        !_PyArg_NoKeywords("chain", kwds))
1913
0
        return NULL;
1914
1915
0
    source = PyObject_GetIter(args);
1916
0
    if (source == NULL)
1917
0
        return NULL;
1918
1919
0
    return chain_new_internal(type, source);
1920
0
}
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
0
{
1935
0
    PyObject *source;
1936
1937
0
    source = PyObject_GetIter(arg);
1938
0
    if (source == NULL)
1939
0
        return NULL;
1940
1941
0
    return chain_new_internal(type, source);
1942
0
}
1943
1944
static void
1945
chain_dealloc(PyObject *op)
1946
0
{
1947
0
    chainobject *lz = chainobject_CAST(op);
1948
0
    PyTypeObject *tp = Py_TYPE(lz);
1949
0
    PyObject_GC_UnTrack(lz);
1950
0
    Py_XDECREF(lz->active);
1951
0
    Py_XDECREF(lz->source);
1952
0
    tp->tp_free(lz);
1953
0
    Py_DECREF(tp);
1954
0
}
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
0
{
1969
0
    chainobject *lz = chainobject_CAST(op);
1970
0
    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
0
    while (lz->source != NULL) {
1976
0
        if (lz->active == NULL) {
1977
0
            PyObject *iterable = PyIter_Next(lz->source);
1978
0
            if (iterable == NULL) {
1979
0
                Py_CLEAR(lz->source);
1980
0
                return NULL;            /* no more input sources */
1981
0
            }
1982
0
            lz->active = PyObject_GetIter(iterable);
1983
0
            Py_DECREF(iterable);
1984
0
            if (lz->active == NULL) {
1985
0
                Py_CLEAR(lz->source);
1986
0
                return NULL;            /* input not iterable */
1987
0
            }
1988
0
        }
1989
0
        item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1990
0
        if (item != NULL)
1991
0
            return item;
1992
0
        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
0
        Py_CLEAR(lz->active);
2000
0
    }
2001
    /* Everything had been consumed already. */
2002
0
    return NULL;
2003
0
}
2004
2005
static PyObject *
2006
chain_next(PyObject *op)
2007
0
{
2008
0
    PyObject *result;
2009
0
    Py_BEGIN_CRITICAL_SECTION(op);
2010
0
    result = chain_next_lock_held(op);
2011
0
    Py_END_CRITICAL_SECTION()
2012
0
    return result;
2013
0
}
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
504
#define productobject_CAST(op)  ((productobject *)(op))
2067
2068
static PyObject *
2069
product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2070
99
{
2071
99
    productobject *lz;
2072
99
    Py_ssize_t nargs, npools, repeat=1;
2073
99
    PyObject *pools = NULL;
2074
99
    Py_ssize_t *indices = NULL;
2075
99
    Py_ssize_t i;
2076
2077
99
    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
99
    assert(PyTuple_CheckExact(args));
2096
99
    if (repeat == 0) {
2097
0
        nargs = 0;
2098
99
    } else {
2099
99
        nargs = PyTuple_GET_SIZE(args);
2100
99
        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
99
    }
2105
99
    npools = nargs * repeat;
2106
2107
99
    indices = PyMem_New(Py_ssize_t, npools);
2108
99
    if (indices == NULL) {
2109
0
        PyErr_NoMemory();
2110
0
        goto error;
2111
0
    }
2112
2113
99
    pools = PyTuple_New(npools);
2114
99
    if (pools == NULL)
2115
0
        goto error;
2116
2117
252
    for (i=0; i < nargs ; ++i) {
2118
153
        PyObject *item = PyTuple_GET_ITEM(args, i);
2119
0
        PyObject *pool = PySequence_Tuple(item);
2120
153
        if (pool == NULL)
2121
0
            goto error;
2122
153
        PyTuple_SET_ITEM(pools, i, pool);
2123
153
        indices[i] = 0;
2124
153
    }
2125
99
    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
99
    lz = (productobject *)type->tp_alloc(type, 0);
2134
99
    if (lz == NULL)
2135
0
        goto error;
2136
2137
99
    lz->pools = pools;
2138
99
    lz->indices = indices;
2139
99
    lz->result = NULL;
2140
99
    lz->stopped = 0;
2141
2142
99
    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
99
}
2150
2151
static void
2152
product_dealloc(PyObject *op)
2153
99
{
2154
99
    productobject *lz = productobject_CAST(op);
2155
99
    PyTypeObject *tp = Py_TYPE(lz);
2156
99
    PyObject_GC_UnTrack(lz);
2157
99
    Py_XDECREF(lz->pools);
2158
99
    Py_XDECREF(lz->result);
2159
99
    PyMem_Free(lz->indices);
2160
99
    tp->tp_free(lz);
2161
99
    Py_DECREF(tp);
2162
99
}
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
405
{
2188
405
    productobject *lz = productobject_CAST(op);
2189
405
    PyObject *pool;
2190
405
    PyObject *elem;
2191
405
    PyObject *oldelem;
2192
405
    PyObject *pools = lz->pools;
2193
405
    PyObject *result = lz->result;
2194
405
    Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2195
405
    Py_ssize_t i;
2196
2197
405
    if (lz->stopped)
2198
0
        return NULL;
2199
2200
405
    if (result == NULL) {
2201
        /* On the first pass, return an initial tuple filled with the
2202
           first element from each pool. */
2203
99
        result = PyTuple_New(npools);
2204
99
        if (result == NULL)
2205
0
            goto empty;
2206
99
        lz->result = result;
2207
252
        for (i=0; i < npools; i++) {
2208
153
            pool = PyTuple_GET_ITEM(pools, i);
2209
153
            if (PyTuple_GET_SIZE(pool) == 0)
2210
0
                goto empty;
2211
153
            elem = PyTuple_GET_ITEM(pool, 0);
2212
153
            Py_INCREF(elem);
2213
153
            PyTuple_SET_ITEM(result, i, elem);
2214
153
        }
2215
306
    } else {
2216
306
        Py_ssize_t *indices = lz->indices;
2217
2218
        /* Copy the previous result tuple or re-use it if available */
2219
306
        if (!_PyObject_IsUniquelyReferenced(result)) {
2220
306
            PyObject *old_result = result;
2221
306
            result = PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
2222
306
            if (result == NULL)
2223
0
                goto empty;
2224
306
            lz->result = result;
2225
306
            Py_DECREF(old_result);
2226
306
        }
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
306
        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
513
        for (i=npools-1 ; i >= 0 ; i--) {
2238
414
            pool = PyTuple_GET_ITEM(pools, i);
2239
0
            indices[i]++;
2240
414
            if (indices[i] == PyTuple_GET_SIZE(pool)) {
2241
                /* Roll-over and advance to next pool */
2242
207
                indices[i] = 0;
2243
207
                elem = PyTuple_GET_ITEM(pool, 0);
2244
207
                Py_INCREF(elem);
2245
207
                oldelem = PyTuple_GET_ITEM(result, i);
2246
207
                PyTuple_SET_ITEM(result, i, elem);
2247
207
                Py_DECREF(oldelem);
2248
207
            } else {
2249
                /* No rollover. Just increment and stop here. */
2250
207
                elem = PyTuple_GET_ITEM(pool, indices[i]);
2251
207
                Py_INCREF(elem);
2252
207
                oldelem = PyTuple_GET_ITEM(result, i);
2253
207
                PyTuple_SET_ITEM(result, i, elem);
2254
207
                Py_DECREF(oldelem);
2255
207
                break;
2256
207
            }
2257
414
        }
2258
2259
        /* If i is negative, then the indices have all rolled-over
2260
           and we're done. */
2261
306
        if (i < 0)
2262
99
            goto empty;
2263
306
    }
2264
2265
306
    return Py_NewRef(result);
2266
2267
99
empty:
2268
99
    lz->stopped = 1;
2269
99
    return NULL;
2270
405
}
2271
2272
static PyObject *
2273
product_next(PyObject *op)
2274
405
{
2275
405
    PyObject *result;
2276
405
    Py_BEGIN_CRITICAL_SECTION(op);
2277
405
    result = product_next_lock_held(op);
2278
405
    Py_END_CRITICAL_SECTION()
2279
405
    return result;
2280
405
}
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
243
#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
72
{
2848
72
    permutationsobject *po;
2849
72
    Py_ssize_t n;
2850
72
    Py_ssize_t r;
2851
72
    PyObject *pool = NULL;
2852
72
    Py_ssize_t *indices = NULL;
2853
72
    Py_ssize_t *cycles = NULL;
2854
72
    Py_ssize_t i;
2855
2856
72
    pool = PySequence_Tuple(iterable);
2857
72
    if (pool == NULL)
2858
0
        goto error;
2859
72
    n = PyTuple_GET_SIZE(pool);
2860
2861
72
    r = n;
2862
72
    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
72
    if (r < 0) {
2872
0
        PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2873
0
        goto error;
2874
0
    }
2875
2876
72
    indices = PyMem_New(Py_ssize_t, n);
2877
72
    cycles = PyMem_New(Py_ssize_t, r);
2878
72
    if (indices == NULL || cycles == NULL) {
2879
0
        PyErr_NoMemory();
2880
0
        goto error;
2881
0
    }
2882
2883
171
    for (i=0 ; i<n ; i++)
2884
99
        indices[i] = i;
2885
171
    for (i=0 ; i<r ; i++)
2886
99
        cycles[i] = n - i;
2887
2888
    /* create permutationsobject structure */
2889
72
    po = (permutationsobject *)type->tp_alloc(type, 0);
2890
72
    if (po == NULL)
2891
0
        goto error;
2892
2893
72
    po->pool = pool;
2894
72
    po->indices = indices;
2895
72
    po->cycles = cycles;
2896
72
    po->result = NULL;
2897
72
    po->r = r;
2898
72
    po->stopped = r > n ? 1 : 0;
2899
2900
72
    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
72
}
2910
2911
static void
2912
permutations_dealloc(PyObject *op)
2913
72
{
2914
72
    permutationsobject *po = permutationsobject_CAST(op);
2915
72
    PyTypeObject *tp = Py_TYPE(po);
2916
72
    PyObject_GC_UnTrack(po);
2917
72
    Py_XDECREF(po->pool);
2918
72
    Py_XDECREF(po->result);
2919
72
    PyMem_Free(po->indices);
2920
72
    PyMem_Free(po->cycles);
2921
72
    tp->tp_free(po);
2922
72
    Py_DECREF(tp);
2923
72
}
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
171
{
2948
171
    permutationsobject *po = permutationsobject_CAST(op);
2949
171
    PyObject *elem;
2950
171
    PyObject *oldelem;
2951
171
    PyObject *pool = po->pool;
2952
171
    Py_ssize_t *indices = po->indices;
2953
171
    Py_ssize_t *cycles = po->cycles;
2954
171
    PyObject *result = po->result;
2955
171
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2956
171
    Py_ssize_t r = po->r;
2957
171
    Py_ssize_t i, j, k, index;
2958
2959
171
    if (po->stopped)
2960
0
        return NULL;
2961
2962
171
    if (result == NULL) {
2963
        /* On the first pass, initialize result tuple using the indices */
2964
72
        result = PyTuple_New(r);
2965
72
        if (result == NULL)
2966
0
            goto empty;
2967
72
        po->result = result;
2968
171
        for (i=0; i<r ; i++) {
2969
99
            index = indices[i];
2970
99
            elem = PyTuple_GET_ITEM(pool, index);
2971
99
            Py_INCREF(elem);
2972
99
            PyTuple_SET_ITEM(result, i, elem);
2973
99
        }
2974
99
    } else {
2975
99
        if (n == 0)
2976
0
            goto empty;
2977
2978
        /* Copy the previous result tuple or re-use it if available */
2979
99
        if (!_PyObject_IsUniquelyReferenced(result)) {
2980
99
            PyObject *old_result = result;
2981
99
            result = PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2982
99
            if (result == NULL)
2983
0
                goto empty;
2984
99
            po->result = result;
2985
99
            Py_DECREF(old_result);
2986
99
        }
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
99
        assert(r == 0 || Py_REFCNT(result) == 1);
2994
2995
        /* Decrement rightmost cycle, moving leftward upon zero rollover */
2996
225
        for (i=r-1 ; i>=0 ; i--) {
2997
153
            cycles[i] -= 1;
2998
153
            if (cycles[i] == 0) {
2999
                /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3000
126
                index = indices[i];
3001
153
                for (j=i ; j<n-1 ; j++)
3002
27
                    indices[j] = indices[j+1];
3003
126
                indices[n-1] = index;
3004
126
                cycles[i] = n - i;
3005
126
            } else {
3006
27
                j = cycles[i];
3007
27
                index = indices[i];
3008
27
                indices[i] = indices[n-j];
3009
27
                indices[n-j] = index;
3010
3011
81
                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
54
                    index = indices[k];
3015
54
                    elem = PyTuple_GET_ITEM(pool, index);
3016
54
                    Py_INCREF(elem);
3017
54
                    oldelem = PyTuple_GET_ITEM(result, k);
3018
54
                    PyTuple_SET_ITEM(result, k, elem);
3019
54
                    Py_DECREF(oldelem);
3020
54
                }
3021
27
                break;
3022
27
            }
3023
153
        }
3024
        /* If i is negative, then the cycles have all
3025
           rolled-over and we're done. */
3026
99
        if (i < 0)
3027
72
            goto empty;
3028
99
    }
3029
99
    return Py_NewRef(result);
3030
3031
72
empty:
3032
72
    po->stopped = 1;
3033
72
    return NULL;
3034
171
}
3035
3036
static PyObject *
3037
permutations_next(PyObject *op)
3038
171
{
3039
171
    PyObject *result;
3040
171
    Py_BEGIN_CRITICAL_SECTION(op);
3041
171
    result = permutations_next_lock_held(op);
3042
171
    Py_END_CRITICAL_SECTION()
3043
171
    return result;
3044
171
}
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
0
#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
0
{
3377
0
    PyObject *it;
3378
0
    filterfalseobject *lz;
3379
3380
    /* Get iterator. */
3381
0
    it = PyObject_GetIter(seq);
3382
0
    if (it == NULL)
3383
0
        return NULL;
3384
3385
    /* create filterfalseobject structure */
3386
0
    lz = (filterfalseobject *)type->tp_alloc(type, 0);
3387
0
    if (lz == NULL) {
3388
0
        Py_DECREF(it);
3389
0
        return NULL;
3390
0
    }
3391
0
    lz->func = Py_NewRef(func);
3392
0
    lz->it = it;
3393
3394
0
    return (PyObject *)lz;
3395
0
}
3396
3397
static void
3398
filterfalse_dealloc(PyObject *op)
3399
0
{
3400
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3401
0
    PyTypeObject *tp = Py_TYPE(lz);
3402
0
    PyObject_GC_UnTrack(lz);
3403
0
    Py_XDECREF(lz->func);
3404
0
    Py_XDECREF(lz->it);
3405
0
    tp->tp_free(lz);
3406
0
    Py_DECREF(tp);
3407
0
}
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
0
{
3422
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3423
0
    PyObject *item;
3424
0
    PyObject *it = lz->it;
3425
0
    long ok;
3426
0
    PyObject *(*iternext)(PyObject *);
3427
3428
0
    iternext = *Py_TYPE(it)->tp_iternext;
3429
0
    for (;;) {
3430
0
        item = iternext(it);
3431
0
        if (item == NULL)
3432
0
            return NULL;
3433
3434
0
        if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3435
0
            ok = PyObject_IsTrue(item);
3436
0
        } else {
3437
0
            PyObject *good;
3438
0
            good = PyObject_CallOneArg(lz->func, item);
3439
0
            if (good == NULL) {
3440
0
                Py_DECREF(item);
3441
0
                return NULL;
3442
0
            }
3443
0
            ok = PyObject_IsTrue(good);
3444
0
            Py_DECREF(good);
3445
0
        }
3446
0
        if (ok == 0)
3447
0
            return item;
3448
0
        Py_DECREF(item);
3449
0
        if (ok < 0)
3450
0
            return NULL;
3451
0
    }
3452
0
}
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
0
#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
0
{
3524
0
    countobject *lz;
3525
0
    int fast_mode;
3526
0
    Py_ssize_t cnt = 0;
3527
0
    long step;
3528
3529
0
    if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3530
0
        (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
0
    fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3536
0
                (long_step == NULL || PyLong_Check(long_step));
3537
3538
    /* If not specified, start defaults to 0 */
3539
0
    if (long_cnt != NULL) {
3540
0
        if (fast_mode) {
3541
0
            assert(PyLong_Check(long_cnt));
3542
0
            cnt = PyLong_AsSsize_t(long_cnt);
3543
0
            if (cnt == -1 && PyErr_Occurred()) {
3544
0
                PyErr_Clear();
3545
0
                fast_mode = 0;
3546
0
            }
3547
0
        }
3548
0
    } else {
3549
0
        cnt = 0;
3550
0
        long_cnt = _PyLong_GetZero();
3551
0
    }
3552
0
    Py_INCREF(long_cnt);
3553
3554
    /* If not specified, step defaults to 1 */
3555
0
    if (long_step == NULL) {
3556
0
        long_step = _PyLong_GetOne();
3557
0
    }
3558
0
    Py_INCREF(long_step);
3559
3560
0
    assert(long_cnt != NULL && long_step != NULL);
3561
3562
    /* Fast mode only works when the step is 1 */
3563
0
    if (fast_mode) {
3564
0
        assert(PyLong_Check(long_step));
3565
0
        step = PyLong_AsLong(long_step);
3566
0
        if (step != 1) {
3567
0
            fast_mode = 0;
3568
0
            if (step == -1 && PyErr_Occurred())
3569
0
                PyErr_Clear();
3570
0
        }
3571
0
    }
3572
3573
0
    if (fast_mode)
3574
0
        Py_CLEAR(long_cnt);
3575
0
    else
3576
0
        cnt = PY_SSIZE_T_MAX;
3577
3578
0
    assert((long_cnt == NULL && fast_mode) ||
3579
0
           (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
3580
0
    assert(!fast_mode ||
3581
0
           (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
3582
3583
    /* create countobject structure */
3584
0
    lz = (countobject *)type->tp_alloc(type, 0);
3585
0
    if (lz == NULL) {
3586
0
        Py_XDECREF(long_cnt);
3587
0
        Py_DECREF(long_step);
3588
0
        return NULL;
3589
0
    }
3590
0
    lz->cnt = cnt;
3591
0
    lz->long_cnt = long_cnt;
3592
0
    lz->long_step = long_step;
3593
3594
0
    return (PyObject *)lz;
3595
0
}
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
0
{
3612
0
    countobject *lz = countobject_CAST(op);
3613
0
    Py_VISIT(Py_TYPE(lz));
3614
0
    Py_VISIT(lz->long_cnt);
3615
0
    Py_VISIT(lz->long_step);
3616
0
    return 0;
3617
0
}
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
0
{
3647
0
    countobject *lz = countobject_CAST(op);
3648
0
#ifndef Py_GIL_DISABLED
3649
0
    if (lz->cnt == PY_SSIZE_T_MAX)
3650
0
        return count_nextlong(lz);
3651
0
    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
0
}
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
        Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(lz->cnt);
3680
0
        return PyUnicode_FromFormat("%s(%zd)",
3681
0
                                    _PyType_Name(Py_TYPE(lz)), cnt);
3682
0
    }
3683
3684
0
    if (PyLong_Check(lz->long_step)) {
3685
0
        long step = PyLong_AsLong(lz->long_step);
3686
0
        if (step == -1 && PyErr_Occurred()) {
3687
0
            PyErr_Clear();
3688
0
        }
3689
0
        if (step == 1) {
3690
            /* Don't display step when it is an integer equal to 1 */
3691
0
            return PyUnicode_FromFormat("%s(%R)",
3692
0
                                        _PyType_Name(Py_TYPE(lz)),
3693
0
                                        lz->long_cnt);
3694
0
        }
3695
0
    }
3696
0
    return PyUnicode_FromFormat("%s(%R, %R)",
3697
0
                                _PyType_Name(Py_TYPE(lz)),
3698
0
                                lz->long_cnt, lz->long_step);
3699
0
}
3700
3701
static PyType_Slot count_slots[] = {
3702
    {Py_tp_dealloc, count_dealloc},
3703
    {Py_tp_repr, count_repr},
3704
    {Py_tp_getattro, PyObject_GenericGetAttr},
3705
    {Py_tp_doc, (void *)itertools_count__doc__},
3706
    {Py_tp_traverse, count_traverse},
3707
    {Py_tp_iter, PyObject_SelfIter},
3708
    {Py_tp_iternext, count_next},
3709
    {Py_tp_new, itertools_count},
3710
    {Py_tp_free, PyObject_GC_Del},
3711
    {0, NULL},
3712
};
3713
3714
static PyType_Spec count_spec = {
3715
    .name = "itertools.count",
3716
    .basicsize = sizeof(countobject),
3717
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3718
              Py_TPFLAGS_IMMUTABLETYPE),
3719
    .slots = count_slots,
3720
};
3721
3722
3723
/* repeat object ************************************************************/
3724
3725
typedef struct {
3726
    PyObject_HEAD
3727
    PyObject *element;
3728
    Py_ssize_t cnt;
3729
} repeatobject;
3730
3731
0
#define repeatobject_CAST(op)   ((repeatobject *)(op))
3732
3733
static PyObject *
3734
repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3735
0
{
3736
0
    repeatobject *ro;
3737
0
    PyObject *element;
3738
0
    Py_ssize_t cnt = -1, n_args;
3739
0
    static char *kwargs[] = {"object", "times", NULL};
3740
3741
0
    n_args = PyTuple_GET_SIZE(args);
3742
0
    if (kwds != NULL)
3743
0
        n_args += PyDict_GET_SIZE(kwds);
3744
0
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3745
0
                                     &element, &cnt))
3746
0
        return NULL;
3747
    /* Does user supply times argument? */
3748
0
    if (n_args == 2 && cnt < 0)
3749
0
        cnt = 0;
3750
3751
0
    ro = (repeatobject *)type->tp_alloc(type, 0);
3752
0
    if (ro == NULL)
3753
0
        return NULL;
3754
0
    ro->element = Py_NewRef(element);
3755
0
    ro->cnt = cnt;
3756
0
    return (PyObject *)ro;
3757
0
}
3758
3759
static void
3760
repeat_dealloc(PyObject *op)
3761
0
{
3762
0
    repeatobject *ro = repeatobject_CAST(op);
3763
0
    PyTypeObject *tp = Py_TYPE(ro);
3764
0
    PyObject_GC_UnTrack(ro);
3765
0
    Py_XDECREF(ro->element);
3766
0
    tp->tp_free(ro);
3767
0
    Py_DECREF(tp);
3768
0
}
3769
3770
static int
3771
repeat_traverse(PyObject *op, visitproc visit, void *arg)
3772
0
{
3773
0
    repeatobject *ro = repeatobject_CAST(op);
3774
0
    Py_VISIT(Py_TYPE(ro));
3775
0
    Py_VISIT(ro->element);
3776
0
    return 0;
3777
0
}
3778
3779
static PyObject *
3780
repeat_next(PyObject *op)
3781
0
{
3782
0
    repeatobject *ro = repeatobject_CAST(op);
3783
0
    Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(ro->cnt);
3784
0
    if (cnt == 0) {
3785
0
        return NULL;
3786
0
    }
3787
0
    if (cnt > 0) {
3788
0
        cnt--;
3789
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(ro->cnt, cnt);
3790
0
    }
3791
0
    return Py_NewRef(ro->element);
3792
0
}
3793
3794
static PyObject *
3795
repeat_repr(PyObject *op)
3796
0
{
3797
0
    repeatobject *ro = repeatobject_CAST(op);
3798
0
    if (ro->cnt == -1)
3799
0
        return PyUnicode_FromFormat("%s(%R)",
3800
0
                                    _PyType_Name(Py_TYPE(ro)), ro->element);
3801
0
    else
3802
0
        return PyUnicode_FromFormat("%s(%R, %zd)",
3803
0
                                    _PyType_Name(Py_TYPE(ro)), ro->element,
3804
0
                                    ro->cnt);
3805
0
}
3806
3807
static PyObject *
3808
repeat_len(PyObject *op, PyObject *Py_UNUSED(args))
3809
0
{
3810
0
    repeatobject *ro = repeatobject_CAST(op);
3811
0
    if (ro->cnt == -1) {
3812
0
        PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3813
0
        return NULL;
3814
0
    }
3815
0
    return PyLong_FromSize_t(ro->cnt);
3816
0
}
3817
3818
PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
3819
3820
static PyMethodDef repeat_methods[] = {
3821
    {"__length_hint__", repeat_len, METH_NOARGS, length_hint_doc},
3822
    {NULL,              NULL}           /* sentinel */
3823
};
3824
3825
PyDoc_STRVAR(repeat_doc,
3826
"repeat(object [,times]) -> create an iterator which returns the object\n\
3827
for the specified number of times.  If not specified, returns the object\n\
3828
endlessly.");
3829
3830
static PyType_Slot repeat_slots[] = {
3831
    {Py_tp_dealloc, repeat_dealloc},
3832
    {Py_tp_repr, repeat_repr},
3833
    {Py_tp_getattro, PyObject_GenericGetAttr},
3834
    {Py_tp_doc, (void *)repeat_doc},
3835
    {Py_tp_traverse, repeat_traverse},
3836
    {Py_tp_iter, PyObject_SelfIter},
3837
    {Py_tp_iternext, repeat_next},
3838
    {Py_tp_methods, repeat_methods},
3839
    {Py_tp_new, repeat_new},
3840
    {Py_tp_free, PyObject_GC_Del},
3841
    {0, NULL},
3842
};
3843
3844
static PyType_Spec repeat_spec = {
3845
    .name = "itertools.repeat",
3846
    .basicsize = sizeof(repeatobject),
3847
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3848
              Py_TPFLAGS_IMMUTABLETYPE),
3849
    .slots = repeat_slots,
3850
};
3851
3852
3853
/* ziplongest object *********************************************************/
3854
3855
typedef struct {
3856
    PyObject_HEAD
3857
    Py_ssize_t tuplesize;
3858
    Py_ssize_t numactive;
3859
    PyObject *ittuple;                  /* tuple of iterators */
3860
    PyObject *result;
3861
    PyObject *fillvalue;
3862
} ziplongestobject;
3863
3864
0
#define ziplongestobject_CAST(op)   ((ziplongestobject *)(op))
3865
3866
static PyObject *
3867
zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3868
0
{
3869
0
    ziplongestobject *lz;
3870
0
    Py_ssize_t i;
3871
0
    PyObject *ittuple;  /* tuple of iterators */
3872
0
    PyObject *result;
3873
0
    PyObject *fillvalue = Py_None;
3874
0
    Py_ssize_t tuplesize;
3875
3876
0
    if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
3877
0
        fillvalue = NULL;
3878
0
        if (PyDict_GET_SIZE(kwds) == 1) {
3879
0
            fillvalue = PyDict_GetItemWithError(kwds, &_Py_ID(fillvalue));
3880
0
        }
3881
0
        if (fillvalue == NULL) {
3882
0
            if (!PyErr_Occurred()) {
3883
0
                PyErr_SetString(PyExc_TypeError,
3884
0
                    "zip_longest() got an unexpected keyword argument");
3885
0
            }
3886
0
            return NULL;
3887
0
        }
3888
0
    }
3889
3890
    /* args must be a tuple */
3891
0
    assert(PyTuple_Check(args));
3892
0
    tuplesize = PyTuple_GET_SIZE(args);
3893
3894
    /* obtain iterators */
3895
0
    ittuple = PyTuple_New(tuplesize);
3896
0
    if (ittuple == NULL)
3897
0
        return NULL;
3898
0
    for (i=0; i < tuplesize; i++) {
3899
0
        PyObject *item = PyTuple_GET_ITEM(args, i);
3900
0
        PyObject *it = PyObject_GetIter(item);
3901
0
        if (it == NULL) {
3902
0
            Py_DECREF(ittuple);
3903
0
            return NULL;
3904
0
        }
3905
0
        PyTuple_SET_ITEM(ittuple, i, it);
3906
0
    }
3907
3908
    /* create a result holder */
3909
0
    result = PyTuple_New(tuplesize);
3910
0
    if (result == NULL) {
3911
0
        Py_DECREF(ittuple);
3912
0
        return NULL;
3913
0
    }
3914
0
    for (i=0 ; i < tuplesize ; i++) {
3915
0
        Py_INCREF(Py_None);
3916
0
        PyTuple_SET_ITEM(result, i, Py_None);
3917
0
    }
3918
3919
    /* create ziplongestobject structure */
3920
0
    lz = (ziplongestobject *)type->tp_alloc(type, 0);
3921
0
    if (lz == NULL) {
3922
0
        Py_DECREF(ittuple);
3923
0
        Py_DECREF(result);
3924
0
        return NULL;
3925
0
    }
3926
0
    lz->ittuple = ittuple;
3927
0
    lz->tuplesize = tuplesize;
3928
0
    lz->numactive = tuplesize;
3929
0
    lz->result = result;
3930
0
    lz->fillvalue = Py_NewRef(fillvalue);
3931
0
    return (PyObject *)lz;
3932
0
}
3933
3934
static void
3935
zip_longest_dealloc(PyObject *op)
3936
0
{
3937
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3938
0
    PyTypeObject *tp = Py_TYPE(lz);
3939
0
    PyObject_GC_UnTrack(lz);
3940
0
    Py_XDECREF(lz->ittuple);
3941
0
    Py_XDECREF(lz->result);
3942
0
    Py_XDECREF(lz->fillvalue);
3943
0
    tp->tp_free(lz);
3944
0
    Py_DECREF(tp);
3945
0
}
3946
3947
static int
3948
zip_longest_traverse(PyObject *op, visitproc visit, void *arg)
3949
0
{
3950
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3951
0
    Py_VISIT(Py_TYPE(lz));
3952
0
    Py_VISIT(lz->ittuple);
3953
0
    Py_VISIT(lz->result);
3954
0
    Py_VISIT(lz->fillvalue);
3955
0
    return 0;
3956
0
}
3957
3958
static PyObject *
3959
zip_longest_next_lock_held(PyObject *op)
3960
0
{
3961
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3962
0
    Py_ssize_t i;
3963
0
    Py_ssize_t tuplesize = lz->tuplesize;
3964
0
    PyObject *result = lz->result;
3965
0
    PyObject *it;
3966
0
    PyObject *item;
3967
0
    PyObject *olditem;
3968
3969
0
    if (tuplesize == 0)
3970
0
        return NULL;
3971
0
    if (lz->numactive == 0)
3972
0
        return NULL;
3973
0
    if (_PyObject_IsUniquelyReferenced(result)) {
3974
0
        Py_INCREF(result);
3975
0
        for (i=0 ; i < tuplesize ; i++) {
3976
0
            it = PyTuple_GET_ITEM(lz->ittuple, i);
3977
0
            if (it == NULL) {
3978
0
                item = Py_NewRef(lz->fillvalue);
3979
0
            } else {
3980
0
                item = PyIter_Next(it);
3981
0
                if (item == NULL) {
3982
0
                    lz->numactive -= 1;
3983
0
                    if (lz->numactive == 0 || PyErr_Occurred()) {
3984
0
                        lz->numactive = 0;
3985
0
                        Py_DECREF(result);
3986
0
                        return NULL;
3987
0
                    } else {
3988
0
                        item = Py_NewRef(lz->fillvalue);
3989
0
                        PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3990
0
                        Py_DECREF(it);
3991
0
                    }
3992
0
                }
3993
0
            }
3994
0
            olditem = PyTuple_GET_ITEM(result, i);
3995
0
            PyTuple_SET_ITEM(result, i, item);
3996
0
            Py_DECREF(olditem);
3997
0
        }
3998
        // bpo-42536: The GC may have untracked this result tuple. Since we're
3999
        // recycling it, make sure it's tracked again:
4000
0
        _PyTuple_Recycle(result);
4001
0
    } else {
4002
0
        result = PyTuple_New(tuplesize);
4003
0
        if (result == NULL)
4004
0
            return NULL;
4005
0
        for (i=0 ; i < tuplesize ; i++) {
4006
0
            it = PyTuple_GET_ITEM(lz->ittuple, i);
4007
0
            if (it == NULL) {
4008
0
                item = Py_NewRef(lz->fillvalue);
4009
0
            } else {
4010
0
                item = PyIter_Next(it);
4011
0
                if (item == NULL) {
4012
0
                    lz->numactive -= 1;
4013
0
                    if (lz->numactive == 0 || PyErr_Occurred()) {
4014
0
                        lz->numactive = 0;
4015
0
                        Py_DECREF(result);
4016
0
                        return NULL;
4017
0
                    } else {
4018
0
                        item = Py_NewRef(lz->fillvalue);
4019
0
                        PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4020
0
                        Py_DECREF(it);
4021
0
                    }
4022
0
                }
4023
0
            }
4024
0
            PyTuple_SET_ITEM(result, i, item);
4025
0
        }
4026
0
    }
4027
0
    return result;
4028
0
}
4029
4030
static PyObject *
4031
zip_longest_next(PyObject *op)
4032
0
{
4033
0
    PyObject *result;
4034
0
    Py_BEGIN_CRITICAL_SECTION(op);
4035
0
    result = zip_longest_next_lock_held(op);
4036
0
    Py_END_CRITICAL_SECTION()
4037
0
    return result;
4038
0
}
4039
4040
PyDoc_STRVAR(zip_longest_doc,
4041
"zip_longest(*iterables, fillvalue=None)\n\
4042
--\n\
4043
\n\
4044
Return a zip_longest object whose .__next__() method returns a tuple where\n\
4045
the i-th element comes from the i-th iterable argument.  The .__next__()\n\
4046
method continues until the longest iterable in the argument sequence\n\
4047
is exhausted and then it raises StopIteration.  When the shorter iterables\n\
4048
are exhausted, the fillvalue is substituted in their place.  The fillvalue\n\
4049
defaults to None or can be specified by a keyword argument.\n\
4050
");
4051
4052
static PyType_Slot ziplongest_slots[] = {
4053
    {Py_tp_dealloc, zip_longest_dealloc},
4054
    {Py_tp_getattro, PyObject_GenericGetAttr},
4055
    {Py_tp_doc, (void *)zip_longest_doc},
4056
    {Py_tp_traverse, zip_longest_traverse},
4057
    {Py_tp_iter, PyObject_SelfIter},
4058
    {Py_tp_iternext, zip_longest_next},
4059
    {Py_tp_new, zip_longest_new},
4060
    {Py_tp_free, PyObject_GC_Del},
4061
    {0, NULL},
4062
};
4063
4064
static PyType_Spec ziplongest_spec = {
4065
    .name = "itertools.zip_longest",
4066
    .basicsize = sizeof(ziplongestobject),
4067
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
4068
              Py_TPFLAGS_IMMUTABLETYPE),
4069
    .slots = ziplongest_slots,
4070
};
4071
4072
4073
/* module level code ********************************************************/
4074
4075
PyDoc_STRVAR(module_doc,
4076
"Functional tools for creating and using iterators.\n\
4077
\n\
4078
Infinite iterators:\n\
4079
count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
4080
cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
4081
repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
4082
\n\
4083
Iterators terminating on the shortest input sequence:\n\
4084
accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
4085
batched(p, n) --> [p0, p1, ..., p_n-1], [p_n, p_n+1, ..., p_2n-1], ...\n\
4086
chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4087
chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
4088
compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4089
dropwhile(predicate, seq) --> seq[n], seq[n+1], starting when predicate fails\n\
4090
groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
4091
filterfalse(predicate, seq) --> elements of seq where predicate(elem) is False\n\
4092
islice(seq, [start,] stop [, step]) --> elements from\n\
4093
       seq[start:stop:step]\n\
4094
pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
4095
starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
4096
tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
4097
takewhile(predicate, seq) --> seq[0], seq[1], until predicate fails\n\
4098
zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
4099
\n\
4100
Combinatoric generators:\n\
4101
product(p, q, ... [repeat=1]) --> cartesian product\n\
4102
permutations(p[, r])\n\
4103
combinations(p, r)\n\
4104
combinations_with_replacement(p, r)\n\
4105
");
4106
4107
static int
4108
itertoolsmodule_traverse(PyObject *mod, visitproc visit, void *arg)
4109
356
{
4110
356
    itertools_state *state = get_module_state(mod);
4111
356
    Py_VISIT(state->accumulate_type);
4112
356
    Py_VISIT(state->batched_type);
4113
356
    Py_VISIT(state->chain_type);
4114
356
    Py_VISIT(state->combinations_type);
4115
356
    Py_VISIT(state->compress_type);
4116
356
    Py_VISIT(state->count_type);
4117
356
    Py_VISIT(state->cwr_type);
4118
356
    Py_VISIT(state->cycle_type);
4119
356
    Py_VISIT(state->dropwhile_type);
4120
356
    Py_VISIT(state->filterfalse_type);
4121
356
    Py_VISIT(state->groupby_type);
4122
356
    Py_VISIT(state->_grouper_type);
4123
356
    Py_VISIT(state->islice_type);
4124
356
    Py_VISIT(state->pairwise_type);
4125
356
    Py_VISIT(state->permutations_type);
4126
356
    Py_VISIT(state->product_type);
4127
356
    Py_VISIT(state->repeat_type);
4128
356
    Py_VISIT(state->starmap_type);
4129
356
    Py_VISIT(state->takewhile_type);
4130
356
    Py_VISIT(state->tee_type);
4131
356
    Py_VISIT(state->teedataobject_type);
4132
356
    Py_VISIT(state->ziplongest_type);
4133
356
    return 0;
4134
356
}
4135
4136
static int
4137
itertoolsmodule_clear(PyObject *mod)
4138
0
{
4139
0
    itertools_state *state = get_module_state(mod);
4140
0
    Py_CLEAR(state->accumulate_type);
4141
0
    Py_CLEAR(state->batched_type);
4142
0
    Py_CLEAR(state->chain_type);
4143
0
    Py_CLEAR(state->combinations_type);
4144
0
    Py_CLEAR(state->compress_type);
4145
0
    Py_CLEAR(state->count_type);
4146
0
    Py_CLEAR(state->cwr_type);
4147
0
    Py_CLEAR(state->cycle_type);
4148
0
    Py_CLEAR(state->dropwhile_type);
4149
0
    Py_CLEAR(state->filterfalse_type);
4150
0
    Py_CLEAR(state->groupby_type);
4151
0
    Py_CLEAR(state->_grouper_type);
4152
0
    Py_CLEAR(state->islice_type);
4153
0
    Py_CLEAR(state->pairwise_type);
4154
0
    Py_CLEAR(state->permutations_type);
4155
0
    Py_CLEAR(state->product_type);
4156
0
    Py_CLEAR(state->repeat_type);
4157
0
    Py_CLEAR(state->starmap_type);
4158
0
    Py_CLEAR(state->takewhile_type);
4159
0
    Py_CLEAR(state->tee_type);
4160
0
    Py_CLEAR(state->teedataobject_type);
4161
0
    Py_CLEAR(state->ziplongest_type);
4162
0
    return 0;
4163
0
}
4164
4165
static void
4166
itertoolsmodule_free(void *mod)
4167
0
{
4168
0
    (void)itertoolsmodule_clear((PyObject *)mod);
4169
0
}
4170
4171
132
#define ADD_TYPE(module, type, spec)                                     \
4172
132
do {                                                                     \
4173
132
    type = (PyTypeObject *)PyType_FromModuleAndSpec(module, spec, NULL); \
4174
132
    if (type == NULL) {                                                  \
4175
0
        return -1;                                                       \
4176
0
    }                                                                    \
4177
132
    if (PyModule_AddType(module, type) < 0) {                            \
4178
0
        return -1;                                                       \
4179
0
    }                                                                    \
4180
132
} while (0)
4181
4182
static int
4183
itertoolsmodule_exec(PyObject *mod)
4184
6
{
4185
6
    itertools_state *state = get_module_state(mod);
4186
6
    ADD_TYPE(mod, state->accumulate_type, &accumulate_spec);
4187
6
    ADD_TYPE(mod, state->batched_type, &batched_spec);
4188
6
    ADD_TYPE(mod, state->chain_type, &chain_spec);
4189
6
    ADD_TYPE(mod, state->combinations_type, &combinations_spec);
4190
6
    ADD_TYPE(mod, state->compress_type, &compress_spec);
4191
6
    ADD_TYPE(mod, state->count_type, &count_spec);
4192
6
    ADD_TYPE(mod, state->cwr_type, &cwr_spec);
4193
6
    ADD_TYPE(mod, state->cycle_type, &cycle_spec);
4194
6
    ADD_TYPE(mod, state->dropwhile_type, &dropwhile_spec);
4195
6
    ADD_TYPE(mod, state->filterfalse_type, &filterfalse_spec);
4196
6
    ADD_TYPE(mod, state->groupby_type, &groupby_spec);
4197
6
    ADD_TYPE(mod, state->_grouper_type, &_grouper_spec);
4198
6
    ADD_TYPE(mod, state->islice_type, &islice_spec);
4199
6
    ADD_TYPE(mod, state->pairwise_type, &pairwise_spec);
4200
6
    ADD_TYPE(mod, state->permutations_type, &permutations_spec);
4201
6
    ADD_TYPE(mod, state->product_type, &product_spec);
4202
6
    ADD_TYPE(mod, state->repeat_type, &repeat_spec);
4203
6
    ADD_TYPE(mod, state->starmap_type, &starmap_spec);
4204
6
    ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
4205
6
    ADD_TYPE(mod, state->tee_type, &tee_spec);
4206
6
    ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
4207
6
    ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
4208
4209
6
    Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
4210
6
    return 0;
4211
6
}
4212
4213
static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4214
    _Py_ABI_SLOT,
4215
    {Py_mod_exec, itertoolsmodule_exec},
4216
    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
4217
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
4218
    {0, NULL}
4219
};
4220
4221
static PyMethodDef module_methods[] = {
4222
    ITERTOOLS_TEE_METHODDEF
4223
    {NULL, NULL} /* sentinel */
4224
};
4225
4226
4227
static struct PyModuleDef itertoolsmodule = {
4228
    .m_base = PyModuleDef_HEAD_INIT,
4229
    .m_name = "itertools",
4230
    .m_doc = module_doc,
4231
    .m_size = sizeof(itertools_state),
4232
    .m_methods = module_methods,
4233
    .m_slots = itertoolsmodule_slots,
4234
    .m_traverse = itertoolsmodule_traverse,
4235
    .m_clear = itertoolsmodule_clear,
4236
    .m_free = itertoolsmodule_free,
4237
};
4238
4239
PyMODINIT_FUNC
4240
PyInit_itertools(void)
4241
6
{
4242
6
    return PyModuleDef_Init(&itertoolsmodule);
4243
6
}