Coverage Report

Created: 2026-10-31 06:21

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