Coverage Report

Created: 2025-08-26 06:26

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