Coverage Report

Created: 2025-07-04 06:49

/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
5.57k
{
45
5.57k
    void *state = _PyModule_GetState(mod);
46
5.57k
    assert(state != NULL);
47
5.57k
    return (itertools_state *)state;
48
5.57k
}
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
@classmethod
1133
itertools.cycle.__new__
1134
    iterable: object
1135
    /
1136
Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
1137
[clinic start generated code]*/
1138
1139
static PyObject *
1140
itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
1141
/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
1142
0
{
1143
0
    PyObject *it;
1144
0
    PyObject *saved;
1145
0
    cycleobject *lz;
1146
1147
    /* Get iterator. */
1148
0
    it = PyObject_GetIter(iterable);
1149
0
    if (it == NULL)
1150
0
        return NULL;
1151
1152
0
    saved = PyList_New(0);
1153
0
    if (saved == NULL) {
1154
0
        Py_DECREF(it);
1155
0
        return NULL;
1156
0
    }
1157
1158
    /* create cycleobject structure */
1159
0
    lz = (cycleobject *)type->tp_alloc(type, 0);
1160
0
    if (lz == NULL) {
1161
0
        Py_DECREF(it);
1162
0
        Py_DECREF(saved);
1163
0
        return NULL;
1164
0
    }
1165
0
    lz->it = it;
1166
0
    lz->saved = saved;
1167
0
    lz->index = -1;
1168
1169
0
    return (PyObject *)lz;
1170
0
}
1171
1172
static void
1173
cycle_dealloc(PyObject *op)
1174
0
{
1175
0
    cycleobject *lz = cycleobject_CAST(op);
1176
0
    PyTypeObject *tp = Py_TYPE(lz);
1177
0
    PyObject_GC_UnTrack(lz);
1178
0
    Py_XDECREF(lz->it);
1179
0
    Py_XDECREF(lz->saved);
1180
0
    tp->tp_free(lz);
1181
0
    Py_DECREF(tp);
1182
0
}
1183
1184
static int
1185
cycle_traverse(PyObject *op, visitproc visit, void *arg)
1186
0
{
1187
0
    cycleobject *lz = cycleobject_CAST(op);
1188
0
    Py_VISIT(Py_TYPE(lz));
1189
0
    Py_VISIT(lz->it);
1190
0
    Py_VISIT(lz->saved);
1191
0
    return 0;
1192
0
}
1193
1194
static PyObject *
1195
cycle_next(PyObject *op)
1196
0
{
1197
0
    cycleobject *lz = cycleobject_CAST(op);
1198
0
    PyObject *item;
1199
1200
0
    Py_ssize_t index = FT_ATOMIC_LOAD_SSIZE_RELAXED(lz->index);
1201
1202
0
    if (index < 0) {
1203
0
        item = PyIter_Next(lz->it);
1204
0
        if (item != NULL) {
1205
0
            if (PyList_Append(lz->saved, item)) {
1206
0
                Py_DECREF(item);
1207
0
                return NULL;
1208
0
            }
1209
0
            return item;
1210
0
        }
1211
        /* Note:  StopIteration is already cleared by PyIter_Next() */
1212
0
        if (PyErr_Occurred())
1213
0
            return NULL;
1214
0
        index = 0;
1215
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(lz->index, 0);
1216
0
#ifndef Py_GIL_DISABLED
1217
0
        Py_CLEAR(lz->it);
1218
0
#endif
1219
0
    }
1220
0
    if (PyList_GET_SIZE(lz->saved) == 0)
1221
0
        return NULL;
1222
0
    item = PyList_GetItemRef(lz->saved, index);
1223
0
    assert(item);
1224
0
    index++;
1225
0
    if (index >= PyList_GET_SIZE(lz->saved)) {
1226
0
        index = 0;
1227
0
    }
1228
0
    FT_ATOMIC_STORE_SSIZE_RELAXED(lz->index, index);
1229
0
    return item;
1230
0
}
1231
1232
static PyType_Slot cycle_slots[] = {
1233
    {Py_tp_dealloc, cycle_dealloc},
1234
    {Py_tp_getattro, PyObject_GenericGetAttr},
1235
    {Py_tp_doc, (void *)itertools_cycle__doc__},
1236
    {Py_tp_traverse, cycle_traverse},
1237
    {Py_tp_iter, PyObject_SelfIter},
1238
    {Py_tp_iternext, cycle_next},
1239
    {Py_tp_new, itertools_cycle},
1240
    {Py_tp_free, PyObject_GC_Del},
1241
    {0, NULL},
1242
};
1243
1244
static PyType_Spec cycle_spec = {
1245
    .name = "itertools.cycle",
1246
    .basicsize = sizeof(cycleobject),
1247
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1248
              Py_TPFLAGS_IMMUTABLETYPE),
1249
    .slots = cycle_slots,
1250
};
1251
1252
1253
/* dropwhile object **********************************************************/
1254
1255
typedef struct {
1256
    PyObject_HEAD
1257
    PyObject *func;
1258
    PyObject *it;
1259
    long start;
1260
} dropwhileobject;
1261
1262
0
#define dropwhileobject_CAST(op)    ((dropwhileobject *)(op))
1263
1264
/*[clinic input]
1265
@classmethod
1266
itertools.dropwhile.__new__
1267
    predicate as func: object
1268
    iterable as seq: object
1269
    /
1270
Drop items from the iterable while predicate(item) is true.
1271
1272
Afterwards, return every element until the iterable is exhausted.
1273
[clinic start generated code]*/
1274
1275
static PyObject *
1276
itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1277
/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
1278
0
{
1279
0
    PyObject *it;
1280
0
    dropwhileobject *lz;
1281
1282
    /* Get iterator. */
1283
0
    it = PyObject_GetIter(seq);
1284
0
    if (it == NULL)
1285
0
        return NULL;
1286
1287
    /* create dropwhileobject structure */
1288
0
    lz = (dropwhileobject *)type->tp_alloc(type, 0);
1289
0
    if (lz == NULL) {
1290
0
        Py_DECREF(it);
1291
0
        return NULL;
1292
0
    }
1293
0
    lz->func = Py_NewRef(func);
1294
0
    lz->it = it;
1295
0
    lz->start = 0;
1296
1297
0
    return (PyObject *)lz;
1298
0
}
1299
1300
static void
1301
dropwhile_dealloc(PyObject *op)
1302
0
{
1303
0
    dropwhileobject *lz = dropwhileobject_CAST(op);
1304
0
    PyTypeObject *tp = Py_TYPE(lz);
1305
0
    PyObject_GC_UnTrack(lz);
1306
0
    Py_XDECREF(lz->func);
1307
0
    Py_XDECREF(lz->it);
1308
0
    tp->tp_free(lz);
1309
0
    Py_DECREF(tp);
1310
0
}
1311
1312
static int
1313
dropwhile_traverse(PyObject *op, visitproc visit, void *arg)
1314
0
{
1315
0
    dropwhileobject *lz = dropwhileobject_CAST(op);
1316
0
    Py_VISIT(Py_TYPE(lz));
1317
0
    Py_VISIT(lz->it);
1318
0
    Py_VISIT(lz->func);
1319
0
    return 0;
1320
0
}
1321
1322
static PyObject *
1323
dropwhile_next(PyObject *op)
1324
0
{
1325
0
    dropwhileobject *lz = dropwhileobject_CAST(op);
1326
0
    PyObject *item, *good;
1327
0
    PyObject *it = lz->it;
1328
0
    long ok;
1329
0
    PyObject *(*iternext)(PyObject *);
1330
1331
0
    iternext = *Py_TYPE(it)->tp_iternext;
1332
0
    for (;;) {
1333
0
        item = iternext(it);
1334
0
        if (item == NULL)
1335
0
            return NULL;
1336
0
        if (lz->start == 1)
1337
0
            return item;
1338
1339
0
        good = PyObject_CallOneArg(lz->func, item);
1340
0
        if (good == NULL) {
1341
0
            Py_DECREF(item);
1342
0
            return NULL;
1343
0
        }
1344
0
        ok = PyObject_IsTrue(good);
1345
0
        Py_DECREF(good);
1346
0
        if (ok == 0) {
1347
0
            lz->start = 1;
1348
0
            return item;
1349
0
        }
1350
0
        Py_DECREF(item);
1351
0
        if (ok < 0)
1352
0
            return NULL;
1353
0
    }
1354
0
}
1355
1356
static PyType_Slot dropwhile_slots[] = {
1357
    {Py_tp_dealloc, dropwhile_dealloc},
1358
    {Py_tp_getattro, PyObject_GenericGetAttr},
1359
    {Py_tp_doc, (void *)itertools_dropwhile__doc__},
1360
    {Py_tp_traverse, dropwhile_traverse},
1361
    {Py_tp_iter, PyObject_SelfIter},
1362
    {Py_tp_iternext, dropwhile_next},
1363
    {Py_tp_new, itertools_dropwhile},
1364
    {Py_tp_free, PyObject_GC_Del},
1365
    {0, NULL},
1366
};
1367
1368
static PyType_Spec dropwhile_spec = {
1369
    .name = "itertools.dropwhile",
1370
    .basicsize = sizeof(dropwhileobject),
1371
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1372
              Py_TPFLAGS_IMMUTABLETYPE),
1373
    .slots = dropwhile_slots,
1374
};
1375
1376
1377
/* takewhile object **********************************************************/
1378
1379
typedef struct {
1380
    PyObject_HEAD
1381
    PyObject *func;
1382
    PyObject *it;
1383
    long stop;
1384
} takewhileobject;
1385
1386
0
#define takewhileobject_CAST(op)    ((takewhileobject *)(op))
1387
1388
/*[clinic input]
1389
@classmethod
1390
itertools.takewhile.__new__
1391
    predicate as func: object
1392
    iterable as seq: object
1393
    /
1394
Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1395
[clinic start generated code]*/
1396
1397
static PyObject *
1398
itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1399
/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
1400
0
{
1401
0
    PyObject *it;
1402
0
    takewhileobject *lz;
1403
1404
    /* Get iterator. */
1405
0
    it = PyObject_GetIter(seq);
1406
0
    if (it == NULL)
1407
0
        return NULL;
1408
1409
    /* create takewhileobject structure */
1410
0
    lz = (takewhileobject *)type->tp_alloc(type, 0);
1411
0
    if (lz == NULL) {
1412
0
        Py_DECREF(it);
1413
0
        return NULL;
1414
0
    }
1415
0
    lz->func = Py_NewRef(func);
1416
0
    lz->it = it;
1417
0
    lz->stop = 0;
1418
1419
0
    return (PyObject *)lz;
1420
0
}
1421
1422
static void
1423
takewhile_dealloc(PyObject *op)
1424
0
{
1425
0
    takewhileobject *lz = takewhileobject_CAST(op);
1426
0
    PyTypeObject *tp = Py_TYPE(lz);
1427
0
    PyObject_GC_UnTrack(lz);
1428
0
    Py_XDECREF(lz->func);
1429
0
    Py_XDECREF(lz->it);
1430
0
    tp->tp_free(lz);
1431
0
    Py_DECREF(tp);
1432
0
}
1433
1434
static int
1435
takewhile_traverse(PyObject *op, visitproc visit, void *arg)
1436
0
{
1437
0
    takewhileobject *lz = takewhileobject_CAST(op);
1438
0
    Py_VISIT(Py_TYPE(lz));
1439
0
    Py_VISIT(lz->it);
1440
0
    Py_VISIT(lz->func);
1441
0
    return 0;
1442
0
}
1443
1444
static PyObject *
1445
takewhile_next(PyObject *op)
1446
0
{
1447
0
    takewhileobject *lz = takewhileobject_CAST(op);
1448
0
    PyObject *item, *good;
1449
0
    PyObject *it = lz->it;
1450
0
    long ok;
1451
1452
0
    if (lz->stop == 1)
1453
0
        return NULL;
1454
1455
0
    item = (*Py_TYPE(it)->tp_iternext)(it);
1456
0
    if (item == NULL)
1457
0
        return NULL;
1458
1459
0
    good = PyObject_CallOneArg(lz->func, item);
1460
0
    if (good == NULL) {
1461
0
        Py_DECREF(item);
1462
0
        return NULL;
1463
0
    }
1464
0
    ok = PyObject_IsTrue(good);
1465
0
    Py_DECREF(good);
1466
0
    if (ok > 0)
1467
0
        return item;
1468
0
    Py_DECREF(item);
1469
0
    if (ok == 0)
1470
0
        lz->stop = 1;
1471
0
    return NULL;
1472
0
}
1473
1474
static PyType_Slot takewhile_slots[] = {
1475
    {Py_tp_dealloc, takewhile_dealloc},
1476
    {Py_tp_getattro, PyObject_GenericGetAttr},
1477
    {Py_tp_doc, (void *)itertools_takewhile__doc__},
1478
    {Py_tp_traverse, takewhile_traverse},
1479
    {Py_tp_iter, PyObject_SelfIter},
1480
    {Py_tp_iternext, takewhile_next},
1481
    {Py_tp_new, itertools_takewhile},
1482
    {Py_tp_free, PyObject_GC_Del},
1483
    {0, NULL},
1484
};
1485
1486
static PyType_Spec takewhile_spec = {
1487
    .name = "itertools.takewhile",
1488
    .basicsize = sizeof(takewhileobject),
1489
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1490
              Py_TPFLAGS_IMMUTABLETYPE),
1491
    .slots = takewhile_slots,
1492
};
1493
1494
1495
/* islice object *************************************************************/
1496
1497
typedef struct {
1498
    PyObject_HEAD
1499
    PyObject *it;
1500
    Py_ssize_t next;
1501
    Py_ssize_t stop;
1502
    Py_ssize_t step;
1503
    Py_ssize_t cnt;
1504
} isliceobject;
1505
1506
0
#define isliceobject_CAST(op)   ((isliceobject *)(op))
1507
1508
static PyObject *
1509
islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1510
0
{
1511
0
    PyObject *seq;
1512
0
    Py_ssize_t start=0, stop=-1, step=1;
1513
0
    PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1514
0
    Py_ssize_t numargs;
1515
0
    isliceobject *lz;
1516
1517
0
    itertools_state *st = find_state_by_type(type);
1518
0
    PyTypeObject *islice_type = st->islice_type;
1519
0
    if ((type == islice_type || type->tp_init == islice_type->tp_init) &&
1520
0
        !_PyArg_NoKeywords("islice", kwds))
1521
0
        return NULL;
1522
1523
0
    if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1524
0
        return NULL;
1525
1526
0
    numargs = PyTuple_Size(args);
1527
0
    if (numargs == 2) {
1528
0
        if (a1 != Py_None) {
1529
0
            stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
1530
0
            if (stop == -1) {
1531
0
                if (PyErr_Occurred())
1532
0
                    PyErr_Clear();
1533
0
                PyErr_SetString(PyExc_ValueError,
1534
0
                   "Stop argument for islice() must be None or "
1535
0
                   "an integer: 0 <= x <= sys.maxsize.");
1536
0
                return NULL;
1537
0
            }
1538
0
        }
1539
0
    } else {
1540
0
        if (a1 != Py_None)
1541
0
            start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
1542
0
        if (start == -1 && PyErr_Occurred())
1543
0
            PyErr_Clear();
1544
0
        if (a2 != Py_None) {
1545
0
            stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
1546
0
            if (stop == -1) {
1547
0
                if (PyErr_Occurred())
1548
0
                    PyErr_Clear();
1549
0
                PyErr_SetString(PyExc_ValueError,
1550
0
                   "Stop argument for islice() must be None or "
1551
0
                   "an integer: 0 <= x <= sys.maxsize.");
1552
0
                return NULL;
1553
0
            }
1554
0
        }
1555
0
    }
1556
0
    if (start<0 || stop<-1) {
1557
0
        PyErr_SetString(PyExc_ValueError,
1558
0
           "Indices for islice() must be None or "
1559
0
           "an integer: 0 <= x <= sys.maxsize.");
1560
0
        return NULL;
1561
0
    }
1562
1563
0
    if (a3 != NULL) {
1564
0
        if (a3 != Py_None)
1565
0
            step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
1566
0
        if (step == -1 && PyErr_Occurred())
1567
0
            PyErr_Clear();
1568
0
    }
1569
0
    if (step<1) {
1570
0
        PyErr_SetString(PyExc_ValueError,
1571
0
           "Step for islice() must be a positive integer or None.");
1572
0
        return NULL;
1573
0
    }
1574
1575
    /* Get iterator. */
1576
0
    it = PyObject_GetIter(seq);
1577
0
    if (it == NULL)
1578
0
        return NULL;
1579
1580
    /* create isliceobject structure */
1581
0
    lz = (isliceobject *)type->tp_alloc(type, 0);
1582
0
    if (lz == NULL) {
1583
0
        Py_DECREF(it);
1584
0
        return NULL;
1585
0
    }
1586
0
    lz->it = it;
1587
0
    lz->next = start;
1588
0
    lz->stop = stop;
1589
0
    lz->step = step;
1590
0
    lz->cnt = 0L;
1591
1592
0
    return (PyObject *)lz;
1593
0
}
1594
1595
static void
1596
islice_dealloc(PyObject *op)
1597
0
{
1598
0
    isliceobject *lz = isliceobject_CAST(op);
1599
0
    PyTypeObject *tp = Py_TYPE(lz);
1600
0
    PyObject_GC_UnTrack(lz);
1601
0
    Py_XDECREF(lz->it);
1602
0
    tp->tp_free(lz);
1603
0
    Py_DECREF(tp);
1604
0
}
1605
1606
static int
1607
islice_traverse(PyObject *op, visitproc visit, void *arg)
1608
0
{
1609
0
    isliceobject *lz = isliceobject_CAST(op);
1610
0
    Py_VISIT(Py_TYPE(lz));
1611
0
    Py_VISIT(lz->it);
1612
0
    return 0;
1613
0
}
1614
1615
static PyObject *
1616
islice_next(PyObject *op)
1617
0
{
1618
0
    isliceobject *lz = isliceobject_CAST(op);
1619
0
    PyObject *item;
1620
0
    PyObject *it = lz->it;
1621
0
    Py_ssize_t stop = lz->stop;
1622
0
    Py_ssize_t oldnext;
1623
0
    PyObject *(*iternext)(PyObject *);
1624
1625
0
    if (it == NULL)
1626
0
        return NULL;
1627
1628
0
    iternext = *Py_TYPE(it)->tp_iternext;
1629
0
    while (lz->cnt < lz->next) {
1630
0
        item = iternext(it);
1631
0
        if (item == NULL)
1632
0
            goto empty;
1633
0
        Py_DECREF(item);
1634
0
        lz->cnt++;
1635
0
    }
1636
0
    if (stop != -1 && lz->cnt >= stop)
1637
0
        goto empty;
1638
0
    item = iternext(it);
1639
0
    if (item == NULL)
1640
0
        goto empty;
1641
0
    lz->cnt++;
1642
0
    oldnext = lz->next;
1643
    /* The (size_t) cast below avoids the danger of undefined
1644
       behaviour from signed integer overflow. */
1645
0
    lz->next += (size_t)lz->step;
1646
0
    if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1647
0
        lz->next = stop;
1648
0
    return item;
1649
1650
0
empty:
1651
0
    Py_CLEAR(lz->it);
1652
0
    return NULL;
1653
0
}
1654
1655
PyDoc_STRVAR(islice_doc,
1656
"islice(iterable, stop) --> islice object\n\
1657
islice(iterable, start, stop[, step]) --> islice object\n\
1658
\n\
1659
Return an iterator whose next() method returns selected values from an\n\
1660
iterable.  If start is specified, will skip all preceding elements;\n\
1661
otherwise, start defaults to zero.  Step defaults to one.  If\n\
1662
specified as another value, step determines how many values are\n\
1663
skipped between successive calls.  Works like a slice() on a list\n\
1664
but returns an iterator.");
1665
1666
static PyType_Slot islice_slots[] = {
1667
    {Py_tp_dealloc, islice_dealloc},
1668
    {Py_tp_getattro, PyObject_GenericGetAttr},
1669
    {Py_tp_doc, (void *)islice_doc},
1670
    {Py_tp_traverse, islice_traverse},
1671
    {Py_tp_iter, PyObject_SelfIter},
1672
    {Py_tp_iternext, islice_next},
1673
    {Py_tp_new, islice_new},
1674
    {Py_tp_free, PyObject_GC_Del},
1675
    {0, NULL},
1676
};
1677
1678
static PyType_Spec islice_spec = {
1679
    .name = "itertools.islice",
1680
    .basicsize = sizeof(isliceobject),
1681
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1682
              Py_TPFLAGS_IMMUTABLETYPE),
1683
    .slots = islice_slots,
1684
};
1685
1686
1687
/* starmap object ************************************************************/
1688
1689
typedef struct {
1690
    PyObject_HEAD
1691
    PyObject *func;
1692
    PyObject *it;
1693
} starmapobject;
1694
1695
0
#define starmapobject_CAST(op)  ((starmapobject *)(op))
1696
1697
/*[clinic input]
1698
@classmethod
1699
itertools.starmap.__new__
1700
    function as func: object
1701
    iterable as seq: object
1702
    /
1703
Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1704
[clinic start generated code]*/
1705
1706
static PyObject *
1707
itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1708
/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
1709
0
{
1710
0
    PyObject *it;
1711
0
    starmapobject *lz;
1712
1713
    /* Get iterator. */
1714
0
    it = PyObject_GetIter(seq);
1715
0
    if (it == NULL)
1716
0
        return NULL;
1717
1718
    /* create starmapobject structure */
1719
0
    lz = (starmapobject *)type->tp_alloc(type, 0);
1720
0
    if (lz == NULL) {
1721
0
        Py_DECREF(it);
1722
0
        return NULL;
1723
0
    }
1724
0
    lz->func = Py_NewRef(func);
1725
0
    lz->it = it;
1726
1727
0
    return (PyObject *)lz;
1728
0
}
1729
1730
static void
1731
starmap_dealloc(PyObject *op)
1732
0
{
1733
0
    starmapobject *lz = starmapobject_CAST(op);
1734
0
    PyTypeObject *tp = Py_TYPE(lz);
1735
0
    PyObject_GC_UnTrack(lz);
1736
0
    Py_XDECREF(lz->func);
1737
0
    Py_XDECREF(lz->it);
1738
0
    tp->tp_free(lz);
1739
0
    Py_DECREF(tp);
1740
0
}
1741
1742
static int
1743
starmap_traverse(PyObject *op, visitproc visit, void *arg)
1744
0
{
1745
0
    starmapobject *lz = starmapobject_CAST(op);
1746
0
    Py_VISIT(Py_TYPE(lz));
1747
0
    Py_VISIT(lz->it);
1748
0
    Py_VISIT(lz->func);
1749
0
    return 0;
1750
0
}
1751
1752
static PyObject *
1753
starmap_next(PyObject *op)
1754
0
{
1755
0
    starmapobject *lz = starmapobject_CAST(op);
1756
0
    PyObject *args;
1757
0
    PyObject *result;
1758
0
    PyObject *it = lz->it;
1759
1760
0
    args = (*Py_TYPE(it)->tp_iternext)(it);
1761
0
    if (args == NULL)
1762
0
        return NULL;
1763
0
    if (!PyTuple_CheckExact(args)) {
1764
0
        PyObject *newargs = PySequence_Tuple(args);
1765
0
        Py_DECREF(args);
1766
0
        if (newargs == NULL)
1767
0
            return NULL;
1768
0
        args = newargs;
1769
0
    }
1770
0
    result = PyObject_Call(lz->func, args, NULL);
1771
0
    Py_DECREF(args);
1772
0
    return result;
1773
0
}
1774
1775
static PyType_Slot starmap_slots[] = {
1776
    {Py_tp_dealloc, starmap_dealloc},
1777
    {Py_tp_getattro, PyObject_GenericGetAttr},
1778
    {Py_tp_doc, (void *)itertools_starmap__doc__},
1779
    {Py_tp_traverse, starmap_traverse},
1780
    {Py_tp_iter, PyObject_SelfIter},
1781
    {Py_tp_iternext, starmap_next},
1782
    {Py_tp_new, itertools_starmap},
1783
    {Py_tp_free, PyObject_GC_Del},
1784
    {0, NULL},
1785
};
1786
1787
static PyType_Spec starmap_spec = {
1788
    .name = "itertools.starmap",
1789
    .basicsize = sizeof(starmapobject),
1790
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1791
              Py_TPFLAGS_IMMUTABLETYPE),
1792
    .slots = starmap_slots,
1793
};
1794
1795
1796
/* chain object **************************************************************/
1797
1798
typedef struct {
1799
    PyObject_HEAD
1800
    PyObject *source;                   /* Iterator over input iterables */
1801
    PyObject *active;                   /* Currently running input iterator */
1802
} chainobject;
1803
1804
0
#define chainobject_CAST(op)    ((chainobject *)(op))
1805
1806
static PyObject *
1807
chain_new_internal(PyTypeObject *type, PyObject *source)
1808
0
{
1809
0
    chainobject *lz;
1810
1811
0
    lz = (chainobject *)type->tp_alloc(type, 0);
1812
0
    if (lz == NULL) {
1813
0
        Py_DECREF(source);
1814
0
        return NULL;
1815
0
    }
1816
1817
0
    lz->source = source;
1818
0
    lz->active = NULL;
1819
0
    return (PyObject *)lz;
1820
0
}
1821
1822
static PyObject *
1823
chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1824
0
{
1825
0
    PyObject *source;
1826
1827
0
    itertools_state *state = find_state_by_type(type);
1828
0
    PyTypeObject *chain_type = state->chain_type;
1829
0
    if ((type == chain_type || type->tp_init == chain_type->tp_init) &&
1830
0
        !_PyArg_NoKeywords("chain", kwds))
1831
0
        return NULL;
1832
1833
0
    source = PyObject_GetIter(args);
1834
0
    if (source == NULL)
1835
0
        return NULL;
1836
1837
0
    return chain_new_internal(type, source);
1838
0
}
1839
1840
/*[clinic input]
1841
@classmethod
1842
itertools.chain.from_iterable
1843
    iterable as arg: object
1844
    /
1845
Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1846
[clinic start generated code]*/
1847
1848
static PyObject *
1849
itertools_chain_from_iterable_impl(PyTypeObject *type, PyObject *arg)
1850
/*[clinic end generated code: output=3d7ea7d46b9e43f5 input=72c39e3a2ca3be85]*/
1851
0
{
1852
0
    PyObject *source;
1853
1854
0
    source = PyObject_GetIter(arg);
1855
0
    if (source == NULL)
1856
0
        return NULL;
1857
1858
0
    return chain_new_internal(type, source);
1859
0
}
1860
1861
static void
1862
chain_dealloc(PyObject *op)
1863
0
{
1864
0
    chainobject *lz = chainobject_CAST(op);
1865
0
    PyTypeObject *tp = Py_TYPE(lz);
1866
0
    PyObject_GC_UnTrack(lz);
1867
0
    Py_XDECREF(lz->active);
1868
0
    Py_XDECREF(lz->source);
1869
0
    tp->tp_free(lz);
1870
0
    Py_DECREF(tp);
1871
0
}
1872
1873
static int
1874
chain_traverse(PyObject *op, visitproc visit, void *arg)
1875
0
{
1876
0
    chainobject *lz = chainobject_CAST(op);
1877
0
    Py_VISIT(Py_TYPE(lz));
1878
0
    Py_VISIT(lz->source);
1879
0
    Py_VISIT(lz->active);
1880
0
    return 0;
1881
0
}
1882
1883
static inline PyObject *
1884
chain_next_lock_held(PyObject *op)
1885
0
{
1886
0
    chainobject *lz = chainobject_CAST(op);
1887
0
    PyObject *item;
1888
1889
    /* lz->source is the iterator of iterables. If it's NULL, we've already
1890
     * consumed them all. lz->active is the current iterator. If it's NULL,
1891
     * we should grab a new one from lz->source. */
1892
0
    while (lz->source != NULL) {
1893
0
        if (lz->active == NULL) {
1894
0
            PyObject *iterable = PyIter_Next(lz->source);
1895
0
            if (iterable == NULL) {
1896
0
                Py_CLEAR(lz->source);
1897
0
                return NULL;            /* no more input sources */
1898
0
            }
1899
0
            lz->active = PyObject_GetIter(iterable);
1900
0
            Py_DECREF(iterable);
1901
0
            if (lz->active == NULL) {
1902
0
                Py_CLEAR(lz->source);
1903
0
                return NULL;            /* input not iterable */
1904
0
            }
1905
0
        }
1906
0
        item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1907
0
        if (item != NULL)
1908
0
            return item;
1909
0
        if (PyErr_Occurred()) {
1910
0
            if (PyErr_ExceptionMatches(PyExc_StopIteration))
1911
0
                PyErr_Clear();
1912
0
            else
1913
0
                return NULL;            /* input raised an exception */
1914
0
        }
1915
        /* lz->active is consumed, try with the next iterable. */
1916
0
        Py_CLEAR(lz->active);
1917
0
    }
1918
    /* Everything had been consumed already. */
1919
0
    return NULL;
1920
0
}
1921
1922
static PyObject *
1923
chain_next(PyObject *op)
1924
0
{
1925
0
    PyObject *result;
1926
0
    Py_BEGIN_CRITICAL_SECTION(op);
1927
0
    result = chain_next_lock_held(op);
1928
0
    Py_END_CRITICAL_SECTION()
1929
0
    return result;
1930
0
}
1931
1932
PyDoc_STRVAR(chain_doc,
1933
"chain(*iterables)\n\
1934
--\n\
1935
\n\
1936
Return a chain object whose .__next__() method returns elements from the\n\
1937
first iterable until it is exhausted, then elements from the next\n\
1938
iterable, until all of the iterables are exhausted.");
1939
1940
static PyMethodDef chain_methods[] = {
1941
    ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
1942
    {"__class_getitem__",    Py_GenericAlias,
1943
    METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},
1944
    {NULL,              NULL}           /* sentinel */
1945
};
1946
1947
static PyType_Slot chain_slots[] = {
1948
    {Py_tp_dealloc, chain_dealloc},
1949
    {Py_tp_getattro, PyObject_GenericGetAttr},
1950
    {Py_tp_doc, (void *)chain_doc},
1951
    {Py_tp_traverse, chain_traverse},
1952
    {Py_tp_iter, PyObject_SelfIter},
1953
    {Py_tp_iternext, chain_next},
1954
    {Py_tp_methods, chain_methods},
1955
    {Py_tp_new, chain_new},
1956
    {Py_tp_free, PyObject_GC_Del},
1957
    {0, NULL},
1958
};
1959
1960
static PyType_Spec chain_spec = {
1961
    .name = "itertools.chain",
1962
    .basicsize = sizeof(chainobject),
1963
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1964
              Py_TPFLAGS_IMMUTABLETYPE),
1965
    .slots = chain_slots,
1966
};
1967
1968
1969
/* product object ************************************************************/
1970
1971
typedef struct {
1972
    PyObject_HEAD
1973
    PyObject *pools;        /* tuple of pool tuples */
1974
    Py_ssize_t *indices;    /* one index per pool */
1975
    PyObject *result;       /* most recently returned result tuple */
1976
    int stopped;            /* set to 1 when the iterator is exhausted */
1977
} productobject;
1978
1979
0
#define productobject_CAST(op)  ((productobject *)(op))
1980
1981
static PyObject *
1982
product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1983
0
{
1984
0
    productobject *lz;
1985
0
    Py_ssize_t nargs, npools, repeat=1;
1986
0
    PyObject *pools = NULL;
1987
0
    Py_ssize_t *indices = NULL;
1988
0
    Py_ssize_t i;
1989
1990
0
    if (kwds != NULL) {
1991
0
        char *kwlist[] = {"repeat", 0};
1992
0
        PyObject *tmpargs = PyTuple_New(0);
1993
0
        if (tmpargs == NULL)
1994
0
            return NULL;
1995
0
        if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
1996
0
                                         kwlist, &repeat)) {
1997
0
            Py_DECREF(tmpargs);
1998
0
            return NULL;
1999
0
        }
2000
0
        Py_DECREF(tmpargs);
2001
0
        if (repeat < 0) {
2002
0
            PyErr_SetString(PyExc_ValueError,
2003
0
                            "repeat argument cannot be negative");
2004
0
            return NULL;
2005
0
        }
2006
0
    }
2007
2008
0
    assert(PyTuple_CheckExact(args));
2009
0
    if (repeat == 0) {
2010
0
        nargs = 0;
2011
0
    } else {
2012
0
        nargs = PyTuple_GET_SIZE(args);
2013
0
        if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
2014
0
            PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2015
0
            return NULL;
2016
0
        }
2017
0
    }
2018
0
    npools = nargs * repeat;
2019
2020
0
    indices = PyMem_New(Py_ssize_t, npools);
2021
0
    if (indices == NULL) {
2022
0
        PyErr_NoMemory();
2023
0
        goto error;
2024
0
    }
2025
2026
0
    pools = PyTuple_New(npools);
2027
0
    if (pools == NULL)
2028
0
        goto error;
2029
2030
0
    for (i=0; i < nargs ; ++i) {
2031
0
        PyObject *item = PyTuple_GET_ITEM(args, i);
2032
0
        PyObject *pool = PySequence_Tuple(item);
2033
0
        if (pool == NULL)
2034
0
            goto error;
2035
0
        PyTuple_SET_ITEM(pools, i, pool);
2036
0
        indices[i] = 0;
2037
0
    }
2038
0
    for ( ; i < npools; ++i) {
2039
0
        PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2040
0
        Py_INCREF(pool);
2041
0
        PyTuple_SET_ITEM(pools, i, pool);
2042
0
        indices[i] = 0;
2043
0
    }
2044
2045
    /* create productobject structure */
2046
0
    lz = (productobject *)type->tp_alloc(type, 0);
2047
0
    if (lz == NULL)
2048
0
        goto error;
2049
2050
0
    lz->pools = pools;
2051
0
    lz->indices = indices;
2052
0
    lz->result = NULL;
2053
0
    lz->stopped = 0;
2054
2055
0
    return (PyObject *)lz;
2056
2057
0
error:
2058
0
    if (indices != NULL)
2059
0
        PyMem_Free(indices);
2060
0
    Py_XDECREF(pools);
2061
0
    return NULL;
2062
0
}
2063
2064
static void
2065
product_dealloc(PyObject *op)
2066
0
{
2067
0
    productobject *lz = productobject_CAST(op);
2068
0
    PyTypeObject *tp = Py_TYPE(lz);
2069
0
    PyObject_GC_UnTrack(lz);
2070
0
    Py_XDECREF(lz->pools);
2071
0
    Py_XDECREF(lz->result);
2072
0
    PyMem_Free(lz->indices);
2073
0
    tp->tp_free(lz);
2074
0
    Py_DECREF(tp);
2075
0
}
2076
2077
static PyObject *
2078
product_sizeof(PyObject *op, PyObject *Py_UNUSED(ignored))
2079
0
{
2080
0
    productobject *lz = productobject_CAST(op);
2081
0
    size_t res = _PyObject_SIZE(Py_TYPE(lz));
2082
0
    res += (size_t)PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2083
0
    return PyLong_FromSize_t(res);
2084
0
}
2085
2086
PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2087
2088
static int
2089
product_traverse(PyObject *op, visitproc visit, void *arg)
2090
0
{
2091
0
    productobject *lz = productobject_CAST(op);
2092
0
    Py_VISIT(Py_TYPE(lz));
2093
0
    Py_VISIT(lz->pools);
2094
0
    Py_VISIT(lz->result);
2095
0
    return 0;
2096
0
}
2097
2098
static PyObject *
2099
product_next_lock_held(PyObject *op)
2100
0
{
2101
0
    productobject *lz = productobject_CAST(op);
2102
0
    PyObject *pool;
2103
0
    PyObject *elem;
2104
0
    PyObject *oldelem;
2105
0
    PyObject *pools = lz->pools;
2106
0
    PyObject *result = lz->result;
2107
0
    Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2108
0
    Py_ssize_t i;
2109
2110
0
    if (lz->stopped)
2111
0
        return NULL;
2112
2113
0
    if (result == NULL) {
2114
        /* On the first pass, return an initial tuple filled with the
2115
           first element from each pool. */
2116
0
        result = PyTuple_New(npools);
2117
0
        if (result == NULL)
2118
0
            goto empty;
2119
0
        lz->result = result;
2120
0
        for (i=0; i < npools; i++) {
2121
0
            pool = PyTuple_GET_ITEM(pools, i);
2122
0
            if (PyTuple_GET_SIZE(pool) == 0)
2123
0
                goto empty;
2124
0
            elem = PyTuple_GET_ITEM(pool, 0);
2125
0
            Py_INCREF(elem);
2126
0
            PyTuple_SET_ITEM(result, i, elem);
2127
0
        }
2128
0
    } else {
2129
0
        Py_ssize_t *indices = lz->indices;
2130
2131
        /* Copy the previous result tuple or re-use it if available */
2132
0
        if (Py_REFCNT(result) > 1) {
2133
0
            PyObject *old_result = result;
2134
0
            result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
2135
0
            if (result == NULL)
2136
0
                goto empty;
2137
0
            lz->result = result;
2138
0
            Py_DECREF(old_result);
2139
0
        }
2140
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2141
        // recycling it, make sure it's tracked again:
2142
0
        else {
2143
0
            _PyTuple_Recycle(result);
2144
0
        }
2145
        /* Now, we've got the only copy so we can update it in-place */
2146
0
        assert (npools==0 || Py_REFCNT(result) == 1);
2147
2148
        /* Update the pool indices right-to-left.  Only advance to the
2149
           next pool when the previous one rolls-over */
2150
0
        for (i=npools-1 ; i >= 0 ; i--) {
2151
0
            pool = PyTuple_GET_ITEM(pools, i);
2152
0
            indices[i]++;
2153
0
            if (indices[i] == PyTuple_GET_SIZE(pool)) {
2154
                /* Roll-over and advance to next pool */
2155
0
                indices[i] = 0;
2156
0
                elem = PyTuple_GET_ITEM(pool, 0);
2157
0
                Py_INCREF(elem);
2158
0
                oldelem = PyTuple_GET_ITEM(result, i);
2159
0
                PyTuple_SET_ITEM(result, i, elem);
2160
0
                Py_DECREF(oldelem);
2161
0
            } else {
2162
                /* No rollover. Just increment and stop here. */
2163
0
                elem = PyTuple_GET_ITEM(pool, indices[i]);
2164
0
                Py_INCREF(elem);
2165
0
                oldelem = PyTuple_GET_ITEM(result, i);
2166
0
                PyTuple_SET_ITEM(result, i, elem);
2167
0
                Py_DECREF(oldelem);
2168
0
                break;
2169
0
            }
2170
0
        }
2171
2172
        /* If i is negative, then the indices have all rolled-over
2173
           and we're done. */
2174
0
        if (i < 0)
2175
0
            goto empty;
2176
0
    }
2177
2178
0
    return Py_NewRef(result);
2179
2180
0
empty:
2181
0
    lz->stopped = 1;
2182
0
    return NULL;
2183
0
}
2184
2185
static PyObject *
2186
product_next(PyObject *op)
2187
0
{
2188
0
    PyObject *result;
2189
0
    Py_BEGIN_CRITICAL_SECTION(op);
2190
0
    result = product_next_lock_held(op);
2191
0
    Py_END_CRITICAL_SECTION()
2192
0
    return result;
2193
0
}
2194
2195
static PyMethodDef product_methods[] = {
2196
    {"__sizeof__", product_sizeof, METH_NOARGS, sizeof_doc},
2197
    {NULL,              NULL}   /* sentinel */
2198
};
2199
2200
PyDoc_STRVAR(product_doc,
2201
"product(*iterables, repeat=1)\n\
2202
--\n\
2203
\n\
2204
Cartesian product of input iterables.  Equivalent to nested for-loops.\n\n\
2205
For example, product(A, B) returns the same as:  ((x,y) for x in A for y in B).\n\
2206
The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2207
cycle in a manner similar to an odometer (with the rightmost element changing\n\
2208
on every iteration).\n\n\
2209
To compute the product of an iterable with itself, specify the number\n\
2210
of repetitions with the optional repeat keyword argument. For example,\n\
2211
product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
2212
product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2213
product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2214
2215
static PyType_Slot product_slots[] = {
2216
    {Py_tp_dealloc, product_dealloc},
2217
    {Py_tp_getattro, PyObject_GenericGetAttr},
2218
    {Py_tp_doc, (void *)product_doc},
2219
    {Py_tp_traverse, product_traverse},
2220
    {Py_tp_iter, PyObject_SelfIter},
2221
    {Py_tp_iternext, product_next},
2222
    {Py_tp_methods, product_methods},
2223
    {Py_tp_new, product_new},
2224
    {Py_tp_free, PyObject_GC_Del},
2225
    {0, NULL},
2226
};
2227
2228
static PyType_Spec product_spec = {
2229
    .name = "itertools.product",
2230
    .basicsize = sizeof(productobject),
2231
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2232
              Py_TPFLAGS_IMMUTABLETYPE),
2233
    .slots = product_slots,
2234
};
2235
2236
2237
/* combinations object *******************************************************/
2238
2239
typedef struct {
2240
    PyObject_HEAD
2241
    PyObject *pool;         /* input converted to a tuple */
2242
    Py_ssize_t *indices;    /* one index per result element */
2243
    PyObject *result;       /* most recently returned result tuple */
2244
    Py_ssize_t r;           /* size of result tuple */
2245
    int stopped;            /* set to 1 when the iterator is exhausted */
2246
} combinationsobject;
2247
2248
0
#define combinationsobject_CAST(op) ((combinationsobject *)(op))
2249
2250
/*[clinic input]
2251
@classmethod
2252
itertools.combinations.__new__
2253
    iterable: object
2254
    r: Py_ssize_t
2255
Return successive r-length combinations of elements in the iterable.
2256
2257
combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2258
[clinic start generated code]*/
2259
2260
static PyObject *
2261
itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2262
                            Py_ssize_t r)
2263
/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
2264
0
{
2265
0
    combinationsobject *co;
2266
0
    Py_ssize_t n;
2267
0
    PyObject *pool = NULL;
2268
0
    Py_ssize_t *indices = NULL;
2269
0
    Py_ssize_t i;
2270
2271
0
    pool = PySequence_Tuple(iterable);
2272
0
    if (pool == NULL)
2273
0
        goto error;
2274
0
    n = PyTuple_GET_SIZE(pool);
2275
0
    if (r < 0) {
2276
0
        PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2277
0
        goto error;
2278
0
    }
2279
2280
0
    indices = PyMem_New(Py_ssize_t, r);
2281
0
    if (indices == NULL) {
2282
0
        PyErr_NoMemory();
2283
0
        goto error;
2284
0
    }
2285
2286
0
    for (i=0 ; i<r ; i++)
2287
0
        indices[i] = i;
2288
2289
    /* create combinationsobject structure */
2290
0
    co = (combinationsobject *)type->tp_alloc(type, 0);
2291
0
    if (co == NULL)
2292
0
        goto error;
2293
2294
0
    co->pool = pool;
2295
0
    co->indices = indices;
2296
0
    co->result = NULL;
2297
0
    co->r = r;
2298
0
    co->stopped = r > n ? 1 : 0;
2299
2300
0
    return (PyObject *)co;
2301
2302
0
error:
2303
0
    if (indices != NULL)
2304
0
        PyMem_Free(indices);
2305
0
    Py_XDECREF(pool);
2306
0
    return NULL;
2307
0
}
2308
2309
static void
2310
combinations_dealloc(PyObject *op)
2311
0
{
2312
0
    combinationsobject *co = combinationsobject_CAST(op);
2313
0
    PyTypeObject *tp = Py_TYPE(co);
2314
0
    PyObject_GC_UnTrack(co);
2315
0
    Py_XDECREF(co->pool);
2316
0
    Py_XDECREF(co->result);
2317
0
    PyMem_Free(co->indices);
2318
0
    tp->tp_free(co);
2319
0
    Py_DECREF(tp);
2320
0
}
2321
2322
static PyObject *
2323
combinations_sizeof(PyObject *op, PyObject *Py_UNUSED(args))
2324
0
{
2325
0
    combinationsobject *co = combinationsobject_CAST(op);
2326
0
    size_t res = _PyObject_SIZE(Py_TYPE(co));
2327
0
    res += (size_t)co->r * sizeof(Py_ssize_t);
2328
0
    return PyLong_FromSize_t(res);
2329
0
}
2330
2331
static int
2332
combinations_traverse(PyObject *op, visitproc visit, void *arg)
2333
0
{
2334
0
    combinationsobject *co = combinationsobject_CAST(op);
2335
0
    Py_VISIT(Py_TYPE(co));
2336
0
    Py_VISIT(co->pool);
2337
0
    Py_VISIT(co->result);
2338
0
    return 0;
2339
0
}
2340
2341
static PyObject *
2342
combinations_next_lock_held(PyObject *op)
2343
0
{
2344
0
    combinationsobject *co = combinationsobject_CAST(op);
2345
0
    PyObject *elem;
2346
0
    PyObject *oldelem;
2347
0
    PyObject *pool = co->pool;
2348
0
    Py_ssize_t *indices = co->indices;
2349
0
    PyObject *result = co->result;
2350
0
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2351
0
    Py_ssize_t r = co->r;
2352
0
    Py_ssize_t i, j, index;
2353
2354
0
    if (co->stopped)
2355
0
        return NULL;
2356
2357
0
    if (result == NULL) {
2358
        /* On the first pass, initialize result tuple using the indices */
2359
0
        result = PyTuple_New(r);
2360
0
        if (result == NULL)
2361
0
            goto empty;
2362
0
        co->result = result;
2363
0
        for (i=0; i<r ; i++) {
2364
0
            index = indices[i];
2365
0
            elem = PyTuple_GET_ITEM(pool, index);
2366
0
            Py_INCREF(elem);
2367
0
            PyTuple_SET_ITEM(result, i, elem);
2368
0
        }
2369
0
    } else {
2370
        /* Copy the previous result tuple or re-use it if available */
2371
0
        if (Py_REFCNT(result) > 1) {
2372
0
            PyObject *old_result = result;
2373
0
            result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2374
0
            if (result == NULL)
2375
0
                goto empty;
2376
0
            co->result = result;
2377
0
            Py_DECREF(old_result);
2378
0
        }
2379
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2380
        // recycling it, make sure it's tracked again:
2381
0
        else {
2382
0
            _PyTuple_Recycle(result);
2383
0
        }
2384
        /* Now, we've got the only copy so we can update it in-place
2385
         * CPython's empty tuple is a singleton and cached in
2386
         * PyTuple's freelist.
2387
         */
2388
0
        assert(r == 0 || Py_REFCNT(result) == 1);
2389
2390
        /* Scan indices right-to-left until finding one that is not
2391
           at its maximum (i + n - r). */
2392
0
        for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2393
0
            ;
2394
2395
        /* If i is negative, then the indices are all at
2396
           their maximum value and we're done. */
2397
0
        if (i < 0)
2398
0
            goto empty;
2399
2400
        /* Increment the current index which we know is not at its
2401
           maximum.  Then move back to the right setting each index
2402
           to its lowest possible value (one higher than the index
2403
           to its left -- this maintains the sort order invariant). */
2404
0
        indices[i]++;
2405
0
        for (j=i+1 ; j<r ; j++)
2406
0
            indices[j] = indices[j-1] + 1;
2407
2408
        /* Update the result tuple for the new indices
2409
           starting with i, the leftmost index that changed */
2410
0
        for ( ; i<r ; i++) {
2411
0
            index = indices[i];
2412
0
            elem = PyTuple_GET_ITEM(pool, index);
2413
0
            Py_INCREF(elem);
2414
0
            oldelem = PyTuple_GET_ITEM(result, i);
2415
0
            PyTuple_SET_ITEM(result, i, elem);
2416
0
            Py_DECREF(oldelem);
2417
0
        }
2418
0
    }
2419
2420
0
    return Py_NewRef(result);
2421
2422
0
empty:
2423
0
    co->stopped = 1;
2424
0
    return NULL;
2425
0
}
2426
2427
static PyObject *
2428
combinations_next(PyObject *op)
2429
0
{
2430
0
    PyObject *result;
2431
0
    Py_BEGIN_CRITICAL_SECTION(op);
2432
0
    result = combinations_next_lock_held(op);
2433
0
    Py_END_CRITICAL_SECTION()
2434
0
    return result;
2435
0
}
2436
2437
static PyMethodDef combinations_methods[] = {
2438
    {"__sizeof__", combinations_sizeof, METH_NOARGS, sizeof_doc},
2439
    {NULL,              NULL}   /* sentinel */
2440
};
2441
2442
static PyType_Slot combinations_slots[] = {
2443
    {Py_tp_dealloc, combinations_dealloc},
2444
    {Py_tp_getattro, PyObject_GenericGetAttr},
2445
    {Py_tp_doc, (void *)itertools_combinations__doc__},
2446
    {Py_tp_traverse, combinations_traverse},
2447
    {Py_tp_iter, PyObject_SelfIter},
2448
    {Py_tp_iternext, combinations_next},
2449
    {Py_tp_methods, combinations_methods},
2450
    {Py_tp_new, itertools_combinations},
2451
    {Py_tp_free, PyObject_GC_Del},
2452
    {0, NULL},
2453
};
2454
2455
static PyType_Spec combinations_spec = {
2456
    .name = "itertools.combinations",
2457
    .basicsize = sizeof(combinationsobject),
2458
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2459
              Py_TPFLAGS_IMMUTABLETYPE),
2460
    .slots = combinations_slots,
2461
};
2462
2463
2464
/* combinations with replacement object **************************************/
2465
2466
/* Equivalent to:
2467
2468
        def combinations_with_replacement(iterable, r):
2469
            "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2470
            # number items returned:  (n+r-1)! / r! / (n-1)!
2471
            pool = tuple(iterable)
2472
            n = len(pool)
2473
            indices = [0] * r
2474
            yield tuple(pool[i] for i in indices)
2475
            while 1:
2476
                for i in reversed(range(r)):
2477
                    if indices[i] != n - 1:
2478
                        break
2479
                else:
2480
                    return
2481
                indices[i:] = [indices[i] + 1] * (r - i)
2482
                yield tuple(pool[i] for i in indices)
2483
2484
        def combinations_with_replacement2(iterable, r):
2485
            'Alternate version that filters from product()'
2486
            pool = tuple(iterable)
2487
            n = len(pool)
2488
            for indices in product(range(n), repeat=r):
2489
                if sorted(indices) == list(indices):
2490
                    yield tuple(pool[i] for i in indices)
2491
*/
2492
typedef struct {
2493
    PyObject_HEAD
2494
    PyObject *pool;         /* input converted to a tuple */
2495
    Py_ssize_t *indices;    /* one index per result element */
2496
    PyObject *result;       /* most recently returned result tuple */
2497
    Py_ssize_t r;           /* size of result tuple */
2498
    int stopped;            /* set to 1 when the cwr iterator is exhausted */
2499
} cwrobject;
2500
2501
0
#define cwrobject_CAST(op)  ((cwrobject *)(op))
2502
2503
/*[clinic input]
2504
@classmethod
2505
itertools.combinations_with_replacement.__new__
2506
    iterable: object
2507
    r: Py_ssize_t
2508
Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2509
2510
combinations_with_replacement('ABC', 2) --> ('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')
2511
[clinic start generated code]*/
2512
2513
static PyObject *
2514
itertools_combinations_with_replacement_impl(PyTypeObject *type,
2515
                                             PyObject *iterable,
2516
                                             Py_ssize_t r)
2517
/*[clinic end generated code: output=48b26856d4e659ca input=1dc58e82a0878fdc]*/
2518
0
{
2519
0
    cwrobject *co;
2520
0
    Py_ssize_t n;
2521
0
    PyObject *pool = NULL;
2522
0
    Py_ssize_t *indices = NULL;
2523
0
    Py_ssize_t i;
2524
2525
0
    pool = PySequence_Tuple(iterable);
2526
0
    if (pool == NULL)
2527
0
        goto error;
2528
0
    n = PyTuple_GET_SIZE(pool);
2529
0
    if (r < 0) {
2530
0
        PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2531
0
        goto error;
2532
0
    }
2533
2534
0
    indices = PyMem_New(Py_ssize_t, r);
2535
0
    if (indices == NULL) {
2536
0
        PyErr_NoMemory();
2537
0
        goto error;
2538
0
    }
2539
2540
0
    for (i=0 ; i<r ; i++)
2541
0
        indices[i] = 0;
2542
2543
    /* create cwrobject structure */
2544
0
    co = (cwrobject *)type->tp_alloc(type, 0);
2545
0
    if (co == NULL)
2546
0
        goto error;
2547
2548
0
    co->pool = pool;
2549
0
    co->indices = indices;
2550
0
    co->result = NULL;
2551
0
    co->r = r;
2552
0
    co->stopped = !n && r;
2553
2554
0
    return (PyObject *)co;
2555
2556
0
error:
2557
0
    if (indices != NULL)
2558
0
        PyMem_Free(indices);
2559
0
    Py_XDECREF(pool);
2560
0
    return NULL;
2561
0
}
2562
2563
static void
2564
cwr_dealloc(PyObject *op)
2565
0
{
2566
0
    cwrobject *co = cwrobject_CAST(op);
2567
0
    PyTypeObject *tp = Py_TYPE(co);
2568
0
    PyObject_GC_UnTrack(co);
2569
0
    Py_XDECREF(co->pool);
2570
0
    Py_XDECREF(co->result);
2571
0
    PyMem_Free(co->indices);
2572
0
    tp->tp_free(co);
2573
0
    Py_DECREF(tp);
2574
0
}
2575
2576
static PyObject *
2577
cwr_sizeof(PyObject *op, PyObject *Py_UNUSED(args))
2578
0
{
2579
0
    cwrobject *co = cwrobject_CAST(op);
2580
0
    size_t res = _PyObject_SIZE(Py_TYPE(co));
2581
0
    res += (size_t)co->r * sizeof(Py_ssize_t);
2582
0
    return PyLong_FromSize_t(res);
2583
0
}
2584
2585
static int
2586
cwr_traverse(PyObject *op, visitproc visit, void *arg)
2587
0
{
2588
0
    cwrobject *co = cwrobject_CAST(op);
2589
0
    Py_VISIT(Py_TYPE(co));
2590
0
    Py_VISIT(co->pool);
2591
0
    Py_VISIT(co->result);
2592
0
    return 0;
2593
0
}
2594
2595
static PyObject *
2596
cwr_next(PyObject *op)
2597
0
{
2598
0
    cwrobject *co = cwrobject_CAST(op);
2599
0
    PyObject *elem;
2600
0
    PyObject *oldelem;
2601
0
    PyObject *pool = co->pool;
2602
0
    Py_ssize_t *indices = co->indices;
2603
0
    PyObject *result = co->result;
2604
0
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2605
0
    Py_ssize_t r = co->r;
2606
0
    Py_ssize_t i, index;
2607
2608
0
    if (co->stopped)
2609
0
        return NULL;
2610
2611
0
    if (result == NULL) {
2612
        /* On the first pass, initialize result tuple with pool[0] */
2613
0
        result = PyTuple_New(r);
2614
0
        if (result == NULL)
2615
0
            goto empty;
2616
0
        co->result = result;
2617
0
        if (n > 0) {
2618
0
            elem = PyTuple_GET_ITEM(pool, 0);
2619
0
            for (i=0; i<r ; i++) {
2620
0
                assert(indices[i] == 0);
2621
0
                Py_INCREF(elem);
2622
0
                PyTuple_SET_ITEM(result, i, elem);
2623
0
            }
2624
0
        }
2625
0
    } else {
2626
        /* Copy the previous result tuple or re-use it if available */
2627
0
        if (Py_REFCNT(result) > 1) {
2628
0
            PyObject *old_result = result;
2629
0
            result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2630
0
            if (result == NULL)
2631
0
                goto empty;
2632
0
            co->result = result;
2633
0
            Py_DECREF(old_result);
2634
0
        }
2635
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2636
        // recycling it, make sure it's tracked again:
2637
0
        else {
2638
0
            _PyTuple_Recycle(result);
2639
0
        }
2640
        /* Now, we've got the only copy so we can update it in-place CPython's
2641
           empty tuple is a singleton and cached in PyTuple's freelist. */
2642
0
        assert(r == 0 || Py_REFCNT(result) == 1);
2643
2644
       /* Scan indices right-to-left until finding one that is not
2645
        * at its maximum (n-1). */
2646
0
        for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2647
0
            ;
2648
2649
        /* If i is negative, then the indices are all at
2650
           their maximum value and we're done. */
2651
0
        if (i < 0)
2652
0
            goto empty;
2653
2654
        /* Increment the current index which we know is not at its
2655
           maximum.  Then set all to the right to the same value. */
2656
0
        index = indices[i] + 1;
2657
0
        assert(index < n);
2658
0
        elem = PyTuple_GET_ITEM(pool, index);
2659
0
        for ( ; i<r ; i++) {
2660
0
            indices[i] = index;
2661
0
            Py_INCREF(elem);
2662
0
            oldelem = PyTuple_GET_ITEM(result, i);
2663
0
            PyTuple_SET_ITEM(result, i, elem);
2664
0
            Py_DECREF(oldelem);
2665
0
        }
2666
0
    }
2667
2668
0
    return Py_NewRef(result);
2669
2670
0
empty:
2671
0
    co->stopped = 1;
2672
0
    return NULL;
2673
0
}
2674
2675
static PyMethodDef cwr_methods[] = {
2676
    {"__sizeof__", cwr_sizeof, METH_NOARGS, sizeof_doc},
2677
    {NULL,              NULL}   /* sentinel */
2678
};
2679
2680
static PyType_Slot cwr_slots[] = {
2681
    {Py_tp_dealloc, cwr_dealloc},
2682
    {Py_tp_getattro, PyObject_GenericGetAttr},
2683
    {Py_tp_doc, (void *)itertools_combinations_with_replacement__doc__},
2684
    {Py_tp_traverse, cwr_traverse},
2685
    {Py_tp_iter, PyObject_SelfIter},
2686
    {Py_tp_iternext, cwr_next},
2687
    {Py_tp_methods, cwr_methods},
2688
    {Py_tp_new, itertools_combinations_with_replacement},
2689
    {Py_tp_free, PyObject_GC_Del},
2690
    {0, NULL},
2691
};
2692
2693
static PyType_Spec cwr_spec = {
2694
    .name = "itertools.combinations_with_replacement",
2695
    .basicsize = sizeof(cwrobject),
2696
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2697
              Py_TPFLAGS_IMMUTABLETYPE),
2698
    .slots = cwr_slots,
2699
};
2700
2701
2702
/* permutations object ********************************************************
2703
2704
def permutations(iterable, r=None):
2705
    # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
2706
    # permutations(range(3)) --> 012 021 102 120 201 210
2707
    pool = tuple(iterable)
2708
    n = len(pool)
2709
    r = n if r is None else r
2710
    if r > n:
2711
        return
2712
    indices = list(range(n))
2713
    cycles = list(range(n, n-r, -1))
2714
    yield tuple(pool[i] for i in indices[:r])
2715
    while n:
2716
        for i in reversed(range(r)):
2717
            cycles[i] -= 1
2718
            if cycles[i] == 0:
2719
                indices[i:] = indices[i+1:] + indices[i:i+1]
2720
                cycles[i] = n - i
2721
            else:
2722
                j = cycles[i]
2723
                indices[i], indices[-j] = indices[-j], indices[i]
2724
                yield tuple(pool[i] for i in indices[:r])
2725
                break
2726
        else:
2727
            return
2728
*/
2729
2730
typedef struct {
2731
    PyObject_HEAD
2732
    PyObject *pool;         /* input converted to a tuple */
2733
    Py_ssize_t *indices;    /* one index per element in the pool */
2734
    Py_ssize_t *cycles;     /* one rollover counter per element in the result */
2735
    PyObject *result;       /* most recently returned result tuple */
2736
    Py_ssize_t r;           /* size of result tuple */
2737
    int stopped;            /* set to 1 when the iterator is exhausted */
2738
} permutationsobject;
2739
2740
0
#define permutationsobject_CAST(op) ((permutationsobject *)(op))
2741
2742
/*[clinic input]
2743
@classmethod
2744
itertools.permutations.__new__
2745
    iterable: object
2746
    r as robj: object = None
2747
Return successive r-length permutations of elements in the iterable.
2748
2749
permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
2750
[clinic start generated code]*/
2751
2752
static PyObject *
2753
itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
2754
                            PyObject *robj)
2755
/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
2756
0
{
2757
0
    permutationsobject *po;
2758
0
    Py_ssize_t n;
2759
0
    Py_ssize_t r;
2760
0
    PyObject *pool = NULL;
2761
0
    Py_ssize_t *indices = NULL;
2762
0
    Py_ssize_t *cycles = NULL;
2763
0
    Py_ssize_t i;
2764
2765
0
    pool = PySequence_Tuple(iterable);
2766
0
    if (pool == NULL)
2767
0
        goto error;
2768
0
    n = PyTuple_GET_SIZE(pool);
2769
2770
0
    r = n;
2771
0
    if (robj != Py_None) {
2772
0
        if (!PyLong_Check(robj)) {
2773
0
            PyErr_SetString(PyExc_TypeError, "Expected int as r");
2774
0
            goto error;
2775
0
        }
2776
0
        r = PyLong_AsSsize_t(robj);
2777
0
        if (r == -1 && PyErr_Occurred())
2778
0
            goto error;
2779
0
    }
2780
0
    if (r < 0) {
2781
0
        PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2782
0
        goto error;
2783
0
    }
2784
2785
0
    indices = PyMem_New(Py_ssize_t, n);
2786
0
    cycles = PyMem_New(Py_ssize_t, r);
2787
0
    if (indices == NULL || cycles == NULL) {
2788
0
        PyErr_NoMemory();
2789
0
        goto error;
2790
0
    }
2791
2792
0
    for (i=0 ; i<n ; i++)
2793
0
        indices[i] = i;
2794
0
    for (i=0 ; i<r ; i++)
2795
0
        cycles[i] = n - i;
2796
2797
    /* create permutationsobject structure */
2798
0
    po = (permutationsobject *)type->tp_alloc(type, 0);
2799
0
    if (po == NULL)
2800
0
        goto error;
2801
2802
0
    po->pool = pool;
2803
0
    po->indices = indices;
2804
0
    po->cycles = cycles;
2805
0
    po->result = NULL;
2806
0
    po->r = r;
2807
0
    po->stopped = r > n ? 1 : 0;
2808
2809
0
    return (PyObject *)po;
2810
2811
0
error:
2812
0
    if (indices != NULL)
2813
0
        PyMem_Free(indices);
2814
0
    if (cycles != NULL)
2815
0
        PyMem_Free(cycles);
2816
0
    Py_XDECREF(pool);
2817
0
    return NULL;
2818
0
}
2819
2820
static void
2821
permutations_dealloc(PyObject *op)
2822
0
{
2823
0
    permutationsobject *po = permutationsobject_CAST(op);
2824
0
    PyTypeObject *tp = Py_TYPE(po);
2825
0
    PyObject_GC_UnTrack(po);
2826
0
    Py_XDECREF(po->pool);
2827
0
    Py_XDECREF(po->result);
2828
0
    PyMem_Free(po->indices);
2829
0
    PyMem_Free(po->cycles);
2830
0
    tp->tp_free(po);
2831
0
    Py_DECREF(tp);
2832
0
}
2833
2834
static PyObject *
2835
permutations_sizeof(PyObject *op, PyObject *Py_UNUSED(args))
2836
0
{
2837
0
    permutationsobject *po = permutationsobject_CAST(op);
2838
0
    size_t res = _PyObject_SIZE(Py_TYPE(po));
2839
0
    res += (size_t)PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
2840
0
    res += (size_t)po->r * sizeof(Py_ssize_t);
2841
0
    return PyLong_FromSize_t(res);
2842
0
}
2843
2844
static int
2845
permutations_traverse(PyObject *op, visitproc visit, void *arg)
2846
0
{
2847
0
    permutationsobject *po = permutationsobject_CAST(op);
2848
0
    Py_VISIT(Py_TYPE(po));
2849
0
    Py_VISIT(po->pool);
2850
0
    Py_VISIT(po->result);
2851
0
    return 0;
2852
0
}
2853
2854
static PyObject *
2855
permutations_next(PyObject *op)
2856
0
{
2857
0
    permutationsobject *po = permutationsobject_CAST(op);
2858
0
    PyObject *elem;
2859
0
    PyObject *oldelem;
2860
0
    PyObject *pool = po->pool;
2861
0
    Py_ssize_t *indices = po->indices;
2862
0
    Py_ssize_t *cycles = po->cycles;
2863
0
    PyObject *result = po->result;
2864
0
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2865
0
    Py_ssize_t r = po->r;
2866
0
    Py_ssize_t i, j, k, index;
2867
2868
0
    if (po->stopped)
2869
0
        return NULL;
2870
2871
0
    if (result == NULL) {
2872
        /* On the first pass, initialize result tuple using the indices */
2873
0
        result = PyTuple_New(r);
2874
0
        if (result == NULL)
2875
0
            goto empty;
2876
0
        po->result = result;
2877
0
        for (i=0; i<r ; i++) {
2878
0
            index = indices[i];
2879
0
            elem = PyTuple_GET_ITEM(pool, index);
2880
0
            Py_INCREF(elem);
2881
0
            PyTuple_SET_ITEM(result, i, elem);
2882
0
        }
2883
0
    } else {
2884
0
        if (n == 0)
2885
0
            goto empty;
2886
2887
        /* Copy the previous result tuple or re-use it if available */
2888
0
        if (Py_REFCNT(result) > 1) {
2889
0
            PyObject *old_result = result;
2890
0
            result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2891
0
            if (result == NULL)
2892
0
                goto empty;
2893
0
            po->result = result;
2894
0
            Py_DECREF(old_result);
2895
0
        }
2896
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2897
        // recycling it, make sure it's tracked again:
2898
0
        else {
2899
0
            _PyTuple_Recycle(result);
2900
0
        }
2901
        /* Now, we've got the only copy so we can update it in-place */
2902
0
        assert(r == 0 || Py_REFCNT(result) == 1);
2903
2904
        /* Decrement rightmost cycle, moving leftward upon zero rollover */
2905
0
        for (i=r-1 ; i>=0 ; i--) {
2906
0
            cycles[i] -= 1;
2907
0
            if (cycles[i] == 0) {
2908
                /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2909
0
                index = indices[i];
2910
0
                for (j=i ; j<n-1 ; j++)
2911
0
                    indices[j] = indices[j+1];
2912
0
                indices[n-1] = index;
2913
0
                cycles[i] = n - i;
2914
0
            } else {
2915
0
                j = cycles[i];
2916
0
                index = indices[i];
2917
0
                indices[i] = indices[n-j];
2918
0
                indices[n-j] = index;
2919
2920
0
                for (k=i; k<r ; k++) {
2921
                    /* start with i, the leftmost element that changed */
2922
                    /* yield tuple(pool[k] for k in indices[:r]) */
2923
0
                    index = indices[k];
2924
0
                    elem = PyTuple_GET_ITEM(pool, index);
2925
0
                    Py_INCREF(elem);
2926
0
                    oldelem = PyTuple_GET_ITEM(result, k);
2927
0
                    PyTuple_SET_ITEM(result, k, elem);
2928
0
                    Py_DECREF(oldelem);
2929
0
                }
2930
0
                break;
2931
0
            }
2932
0
        }
2933
        /* If i is negative, then the cycles have all
2934
           rolled-over and we're done. */
2935
0
        if (i < 0)
2936
0
            goto empty;
2937
0
    }
2938
0
    return Py_NewRef(result);
2939
2940
0
empty:
2941
0
    po->stopped = 1;
2942
0
    return NULL;
2943
0
}
2944
2945
static PyMethodDef permuations_methods[] = {
2946
    {"__sizeof__", permutations_sizeof, METH_NOARGS, sizeof_doc},
2947
    {NULL,              NULL}   /* sentinel */
2948
};
2949
2950
static PyType_Slot permutations_slots[] = {
2951
    {Py_tp_dealloc, permutations_dealloc},
2952
    {Py_tp_getattro, PyObject_GenericGetAttr},
2953
    {Py_tp_doc, (void *)itertools_permutations__doc__},
2954
    {Py_tp_traverse, permutations_traverse},
2955
    {Py_tp_iter, PyObject_SelfIter},
2956
    {Py_tp_iternext, permutations_next},
2957
    {Py_tp_methods, permuations_methods},
2958
    {Py_tp_new, itertools_permutations},
2959
    {Py_tp_free, PyObject_GC_Del},
2960
    {0, NULL},
2961
};
2962
2963
static PyType_Spec permutations_spec = {
2964
    .name = "itertools.permutations",
2965
    .basicsize = sizeof(permutationsobject),
2966
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2967
              Py_TPFLAGS_IMMUTABLETYPE),
2968
    .slots = permutations_slots,
2969
};
2970
2971
2972
/* accumulate object ********************************************************/
2973
2974
typedef struct {
2975
    PyObject_HEAD
2976
    PyObject *total;
2977
    PyObject *it;
2978
    PyObject *binop;
2979
    PyObject *initial;
2980
    itertools_state *state;
2981
} accumulateobject;
2982
2983
0
#define accumulateobject_CAST(op)   ((accumulateobject *)(op))
2984
2985
/*[clinic input]
2986
@classmethod
2987
itertools.accumulate.__new__
2988
    iterable: object
2989
    func as binop: object = None
2990
    *
2991
    initial: object = None
2992
Return series of accumulated sums (or other binary function results).
2993
[clinic start generated code]*/
2994
2995
static PyObject *
2996
itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
2997
                          PyObject *binop, PyObject *initial)
2998
/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
2999
0
{
3000
0
    PyObject *it;
3001
0
    accumulateobject *lz;
3002
3003
    /* Get iterator. */
3004
0
    it = PyObject_GetIter(iterable);
3005
0
    if (it == NULL)
3006
0
        return NULL;
3007
3008
    /* create accumulateobject structure */
3009
0
    lz = (accumulateobject *)type->tp_alloc(type, 0);
3010
0
    if (lz == NULL) {
3011
0
        Py_DECREF(it);
3012
0
        return NULL;
3013
0
    }
3014
3015
0
    if (binop != Py_None) {
3016
0
        lz->binop = Py_XNewRef(binop);
3017
0
    }
3018
0
    lz->total = NULL;
3019
0
    lz->it = it;
3020
0
    lz->initial = Py_XNewRef(initial);
3021
0
    lz->state = find_state_by_type(type);
3022
0
    return (PyObject *)lz;
3023
0
}
3024
3025
static void
3026
accumulate_dealloc(PyObject *op)
3027
0
{
3028
0
    accumulateobject *lz = accumulateobject_CAST(op);
3029
0
    PyTypeObject *tp = Py_TYPE(lz);
3030
0
    PyObject_GC_UnTrack(lz);
3031
0
    Py_XDECREF(lz->binop);
3032
0
    Py_XDECREF(lz->total);
3033
0
    Py_XDECREF(lz->it);
3034
0
    Py_XDECREF(lz->initial);
3035
0
    tp->tp_free(lz);
3036
0
    Py_DECREF(tp);
3037
0
}
3038
3039
static int
3040
accumulate_traverse(PyObject *op, visitproc visit, void *arg)
3041
0
{
3042
0
    accumulateobject *lz = accumulateobject_CAST(op);
3043
0
    Py_VISIT(Py_TYPE(lz));
3044
0
    Py_VISIT(lz->binop);
3045
0
    Py_VISIT(lz->it);
3046
0
    Py_VISIT(lz->total);
3047
0
    Py_VISIT(lz->initial);
3048
0
    return 0;
3049
0
}
3050
3051
static PyObject *
3052
accumulate_next(PyObject *op)
3053
0
{
3054
0
    accumulateobject *lz = accumulateobject_CAST(op);
3055
0
    PyObject *val, *newtotal;
3056
3057
0
    if (lz->initial != Py_None) {
3058
0
        lz->total = lz->initial;
3059
0
        lz->initial = Py_NewRef(Py_None);
3060
0
        return Py_NewRef(lz->total);
3061
0
    }
3062
0
    val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
3063
0
    if (val == NULL)
3064
0
        return NULL;
3065
3066
0
    if (lz->total == NULL) {
3067
0
        lz->total = Py_NewRef(val);
3068
0
        return lz->total;
3069
0
    }
3070
3071
0
    if (lz->binop == NULL)
3072
0
        newtotal = PyNumber_Add(lz->total, val);
3073
0
    else
3074
0
        newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
3075
0
    Py_DECREF(val);
3076
0
    if (newtotal == NULL)
3077
0
        return NULL;
3078
3079
0
    Py_INCREF(newtotal);
3080
0
    Py_SETREF(lz->total, newtotal);
3081
0
    return newtotal;
3082
0
}
3083
3084
static PyType_Slot accumulate_slots[] = {
3085
    {Py_tp_dealloc, accumulate_dealloc},
3086
    {Py_tp_getattro, PyObject_GenericGetAttr},
3087
    {Py_tp_doc, (void *)itertools_accumulate__doc__},
3088
    {Py_tp_traverse, accumulate_traverse},
3089
    {Py_tp_iter, PyObject_SelfIter},
3090
    {Py_tp_iternext, accumulate_next},
3091
    {Py_tp_new, itertools_accumulate},
3092
    {Py_tp_free, PyObject_GC_Del},
3093
    {0, NULL},
3094
};
3095
3096
static PyType_Spec accumulate_spec = {
3097
    .name = "itertools.accumulate",
3098
    .basicsize = sizeof(accumulateobject),
3099
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3100
              Py_TPFLAGS_IMMUTABLETYPE),
3101
    .slots = accumulate_slots,
3102
};
3103
3104
3105
/* compress object ************************************************************/
3106
3107
/* Equivalent to:
3108
3109
    def compress(data, selectors):
3110
        "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3111
        return (d for d, s in zip(data, selectors) if s)
3112
*/
3113
3114
typedef struct {
3115
    PyObject_HEAD
3116
    PyObject *data;
3117
    PyObject *selectors;
3118
} compressobject;
3119
3120
0
#define compressobject_CAST(op) ((compressobject *)(op))
3121
3122
/*[clinic input]
3123
@classmethod
3124
itertools.compress.__new__
3125
    data as seq1: object
3126
    selectors as seq2: object
3127
Return data elements corresponding to true selector elements.
3128
3129
Forms a shorter iterator from selected data elements using the selectors to
3130
choose the data elements.
3131
[clinic start generated code]*/
3132
3133
static PyObject *
3134
itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3135
/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
3136
0
{
3137
0
    PyObject *data=NULL, *selectors=NULL;
3138
0
    compressobject *lz;
3139
3140
0
    data = PyObject_GetIter(seq1);
3141
0
    if (data == NULL)
3142
0
        goto fail;
3143
0
    selectors = PyObject_GetIter(seq2);
3144
0
    if (selectors == NULL)
3145
0
        goto fail;
3146
3147
    /* create compressobject structure */
3148
0
    lz = (compressobject *)type->tp_alloc(type, 0);
3149
0
    if (lz == NULL)
3150
0
        goto fail;
3151
0
    lz->data = data;
3152
0
    lz->selectors = selectors;
3153
0
    return (PyObject *)lz;
3154
3155
0
fail:
3156
0
    Py_XDECREF(data);
3157
0
    Py_XDECREF(selectors);
3158
0
    return NULL;
3159
0
}
3160
3161
static void
3162
compress_dealloc(PyObject *op)
3163
0
{
3164
0
    compressobject *lz = compressobject_CAST(op);
3165
0
    PyTypeObject *tp = Py_TYPE(lz);
3166
0
    PyObject_GC_UnTrack(lz);
3167
0
    Py_XDECREF(lz->data);
3168
0
    Py_XDECREF(lz->selectors);
3169
0
    tp->tp_free(lz);
3170
0
    Py_DECREF(tp);
3171
0
}
3172
3173
static int
3174
compress_traverse(PyObject *op, visitproc visit, void *arg)
3175
0
{
3176
0
    compressobject *lz = compressobject_CAST(op);
3177
0
    Py_VISIT(Py_TYPE(lz));
3178
0
    Py_VISIT(lz->data);
3179
0
    Py_VISIT(lz->selectors);
3180
0
    return 0;
3181
0
}
3182
3183
static PyObject *
3184
compress_next(PyObject *op)
3185
0
{
3186
0
    compressobject *lz = compressobject_CAST(op);
3187
0
    PyObject *data = lz->data, *selectors = lz->selectors;
3188
0
    PyObject *datum, *selector;
3189
0
    PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3190
0
    PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3191
0
    int ok;
3192
3193
0
    while (1) {
3194
        /* Steps:  get datum, get selector, evaluate selector.
3195
           Order is important (to match the pure python version
3196
           in terms of which input gets a chance to raise an
3197
           exception first).
3198
        */
3199
3200
0
        datum = datanext(data);
3201
0
        if (datum == NULL)
3202
0
            return NULL;
3203
3204
0
        selector = selectornext(selectors);
3205
0
        if (selector == NULL) {
3206
0
            Py_DECREF(datum);
3207
0
            return NULL;
3208
0
        }
3209
3210
0
        ok = PyObject_IsTrue(selector);
3211
0
        Py_DECREF(selector);
3212
0
        if (ok > 0)
3213
0
            return datum;
3214
0
        Py_DECREF(datum);
3215
0
        if (ok < 0)
3216
0
            return NULL;
3217
0
    }
3218
0
}
3219
3220
static PyType_Slot compress_slots[] = {
3221
    {Py_tp_dealloc, compress_dealloc},
3222
    {Py_tp_getattro, PyObject_GenericGetAttr},
3223
    {Py_tp_doc, (void *)itertools_compress__doc__},
3224
    {Py_tp_traverse, compress_traverse},
3225
    {Py_tp_iter, PyObject_SelfIter},
3226
    {Py_tp_iternext, compress_next},
3227
    {Py_tp_new, itertools_compress},
3228
    {Py_tp_free, PyObject_GC_Del},
3229
    {0, NULL},
3230
};
3231
3232
static PyType_Spec compress_spec = {
3233
    .name = "itertools.compress",
3234
    .basicsize = sizeof(compressobject),
3235
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3236
              Py_TPFLAGS_IMMUTABLETYPE),
3237
    .slots = compress_slots,
3238
};
3239
3240
3241
/* filterfalse object ************************************************************/
3242
3243
typedef struct {
3244
    PyObject_HEAD
3245
    PyObject *func;
3246
    PyObject *it;
3247
} filterfalseobject;
3248
3249
0
#define filterfalseobject_CAST(op)  ((filterfalseobject *)(op))
3250
3251
/*[clinic input]
3252
@classmethod
3253
itertools.filterfalse.__new__
3254
    function as func: object
3255
    iterable as seq: object
3256
    /
3257
Return those items of iterable for which function(item) is false.
3258
3259
If function is None, return the items that are false.
3260
[clinic start generated code]*/
3261
3262
static PyObject *
3263
itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3264
/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
3265
0
{
3266
0
    PyObject *it;
3267
0
    filterfalseobject *lz;
3268
3269
    /* Get iterator. */
3270
0
    it = PyObject_GetIter(seq);
3271
0
    if (it == NULL)
3272
0
        return NULL;
3273
3274
    /* create filterfalseobject structure */
3275
0
    lz = (filterfalseobject *)type->tp_alloc(type, 0);
3276
0
    if (lz == NULL) {
3277
0
        Py_DECREF(it);
3278
0
        return NULL;
3279
0
    }
3280
0
    lz->func = Py_NewRef(func);
3281
0
    lz->it = it;
3282
3283
0
    return (PyObject *)lz;
3284
0
}
3285
3286
static void
3287
filterfalse_dealloc(PyObject *op)
3288
0
{
3289
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3290
0
    PyTypeObject *tp = Py_TYPE(lz);
3291
0
    PyObject_GC_UnTrack(lz);
3292
0
    Py_XDECREF(lz->func);
3293
0
    Py_XDECREF(lz->it);
3294
0
    tp->tp_free(lz);
3295
0
    Py_DECREF(tp);
3296
0
}
3297
3298
static int
3299
filterfalse_traverse(PyObject *op, visitproc visit, void *arg)
3300
0
{
3301
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3302
0
    Py_VISIT(Py_TYPE(lz));
3303
0
    Py_VISIT(lz->it);
3304
0
    Py_VISIT(lz->func);
3305
0
    return 0;
3306
0
}
3307
3308
static PyObject *
3309
filterfalse_next(PyObject *op)
3310
0
{
3311
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3312
0
    PyObject *item;
3313
0
    PyObject *it = lz->it;
3314
0
    long ok;
3315
0
    PyObject *(*iternext)(PyObject *);
3316
3317
0
    iternext = *Py_TYPE(it)->tp_iternext;
3318
0
    for (;;) {
3319
0
        item = iternext(it);
3320
0
        if (item == NULL)
3321
0
            return NULL;
3322
3323
0
        if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3324
0
            ok = PyObject_IsTrue(item);
3325
0
        } else {
3326
0
            PyObject *good;
3327
0
            good = PyObject_CallOneArg(lz->func, item);
3328
0
            if (good == NULL) {
3329
0
                Py_DECREF(item);
3330
0
                return NULL;
3331
0
            }
3332
0
            ok = PyObject_IsTrue(good);
3333
0
            Py_DECREF(good);
3334
0
        }
3335
0
        if (ok == 0)
3336
0
            return item;
3337
0
        Py_DECREF(item);
3338
0
        if (ok < 0)
3339
0
            return NULL;
3340
0
    }
3341
0
}
3342
3343
static PyType_Slot filterfalse_slots[] = {
3344
    {Py_tp_dealloc, filterfalse_dealloc},
3345
    {Py_tp_getattro, PyObject_GenericGetAttr},
3346
    {Py_tp_doc, (void *)itertools_filterfalse__doc__},
3347
    {Py_tp_traverse, filterfalse_traverse},
3348
    {Py_tp_iter, PyObject_SelfIter},
3349
    {Py_tp_iternext, filterfalse_next},
3350
    {Py_tp_new, itertools_filterfalse},
3351
    {Py_tp_free, PyObject_GC_Del},
3352
    {0, NULL},
3353
};
3354
3355
static PyType_Spec filterfalse_spec = {
3356
    .name = "itertools.filterfalse",
3357
    .basicsize = sizeof(filterfalseobject),
3358
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3359
              Py_TPFLAGS_IMMUTABLETYPE),
3360
    .slots = filterfalse_slots,
3361
};
3362
3363
3364
/* count object ************************************************************/
3365
3366
typedef struct {
3367
    PyObject_HEAD
3368
    Py_ssize_t cnt;
3369
    PyObject *long_cnt;
3370
    PyObject *long_step;
3371
} countobject;
3372
3373
3.99k
#define countobject_CAST(op)    ((countobject *)(op))
3374
3375
/* Counting logic and invariants:
3376
3377
fast_mode:  when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3378
3379
    assert(long_cnt == NULL && long_step==PyLong(1));
3380
    Advances with:  cnt += 1
3381
    When count hits PY_SSIZE_T_MAX, switch to slow_mode.
3382
3383
slow_mode:  when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3384
3385
    assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3386
    All counting is done with python objects (no overflows or underflows).
3387
    Advances with:  long_cnt += long_step
3388
    Step may be zero -- effectively a slow version of repeat(cnt).
3389
    Either long_cnt or long_step may be a float, Fraction, or Decimal.
3390
*/
3391
3392
/*[clinic input]
3393
@classmethod
3394
itertools.count.__new__
3395
    start as long_cnt: object(c_default="NULL") = 0
3396
    step as long_step: object(c_default="NULL") = 1
3397
Return a count object whose .__next__() method returns consecutive values.
3398
3399
Equivalent to:
3400
    def count(firstval=0, step=1):
3401
        x = firstval
3402
        while 1:
3403
            yield x
3404
            x += step
3405
[clinic start generated code]*/
3406
3407
static PyObject *
3408
itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
3409
                     PyObject *long_step)
3410
/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
3411
2
{
3412
2
    countobject *lz;
3413
2
    int fast_mode;
3414
2
    Py_ssize_t cnt = 0;
3415
2
    long step;
3416
3417
2
    if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3418
2
        (long_step != NULL && !PyNumber_Check(long_step))) {
3419
0
                    PyErr_SetString(PyExc_TypeError, "a number is required");
3420
0
                    return NULL;
3421
0
    }
3422
3423
2
    fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3424
2
                (long_step == NULL || PyLong_Check(long_step));
3425
3426
    /* If not specified, start defaults to 0 */
3427
2
    if (long_cnt != NULL) {
3428
0
        if (fast_mode) {
3429
0
            assert(PyLong_Check(long_cnt));
3430
0
            cnt = PyLong_AsSsize_t(long_cnt);
3431
0
            if (cnt == -1 && PyErr_Occurred()) {
3432
0
                PyErr_Clear();
3433
0
                fast_mode = 0;
3434
0
            }
3435
0
        }
3436
2
    } else {
3437
2
        cnt = 0;
3438
2
        long_cnt = _PyLong_GetZero();
3439
2
    }
3440
2
    Py_INCREF(long_cnt);
3441
3442
    /* If not specified, step defaults to 1 */
3443
2
    if (long_step == NULL) {
3444
2
        long_step = _PyLong_GetOne();
3445
2
    }
3446
2
    Py_INCREF(long_step);
3447
3448
2
    assert(long_cnt != NULL && long_step != NULL);
3449
3450
    /* Fast mode only works when the step is 1 */
3451
2
    if (fast_mode) {
3452
2
        assert(PyLong_Check(long_step));
3453
2
        step = PyLong_AsLong(long_step);
3454
2
        if (step != 1) {
3455
0
            fast_mode = 0;
3456
0
            if (step == -1 && PyErr_Occurred())
3457
0
                PyErr_Clear();
3458
0
        }
3459
2
    }
3460
3461
2
    if (fast_mode)
3462
2
        Py_CLEAR(long_cnt);
3463
0
    else
3464
0
        cnt = PY_SSIZE_T_MAX;
3465
3466
2
    assert((long_cnt == NULL && fast_mode) ||
3467
2
           (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
3468
2
    assert(!fast_mode ||
3469
2
           (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
3470
3471
    /* create countobject structure */
3472
2
    lz = (countobject *)type->tp_alloc(type, 0);
3473
2
    if (lz == NULL) {
3474
0
        Py_XDECREF(long_cnt);
3475
0
        Py_DECREF(long_step);
3476
0
        return NULL;
3477
0
    }
3478
2
    lz->cnt = cnt;
3479
2
    lz->long_cnt = long_cnt;
3480
2
    lz->long_step = long_step;
3481
3482
2
    return (PyObject *)lz;
3483
2
}
3484
3485
static void
3486
count_dealloc(PyObject *op)
3487
0
{
3488
0
    countobject *lz = countobject_CAST(op);
3489
0
    PyTypeObject *tp = Py_TYPE(lz);
3490
0
    PyObject_GC_UnTrack(lz);
3491
0
    Py_XDECREF(lz->long_cnt);
3492
0
    Py_XDECREF(lz->long_step);
3493
0
    tp->tp_free(lz);
3494
0
    Py_DECREF(tp);
3495
0
}
3496
3497
static int
3498
count_traverse(PyObject *op, visitproc visit, void *arg)
3499
3.99k
{
3500
3.99k
    countobject *lz = countobject_CAST(op);
3501
3.99k
    Py_VISIT(Py_TYPE(lz));
3502
3.99k
    Py_VISIT(lz->long_cnt);
3503
3.99k
    Py_VISIT(lz->long_step);
3504
3.99k
    return 0;
3505
3.99k
}
3506
3507
static PyObject *
3508
count_nextlong(countobject *lz)
3509
0
{
3510
0
    PyObject *long_cnt;
3511
0
    PyObject *stepped_up;
3512
3513
0
    long_cnt = lz->long_cnt;
3514
0
    if (long_cnt == NULL) {
3515
        /* Switch to slow_mode */
3516
0
        long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3517
0
        if (long_cnt == NULL)
3518
0
            return NULL;
3519
0
    }
3520
0
    assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
3521
3522
0
    stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3523
0
    if (stepped_up == NULL)
3524
0
        return NULL;
3525
0
    lz->long_cnt = stepped_up;
3526
0
    return long_cnt;
3527
0
}
3528
3529
static PyObject *
3530
count_next(PyObject *op)
3531
0
{
3532
0
    countobject *lz = countobject_CAST(op);
3533
0
#ifndef Py_GIL_DISABLED
3534
0
    if (lz->cnt == PY_SSIZE_T_MAX)
3535
0
        return count_nextlong(lz);
3536
0
    return PyLong_FromSsize_t(lz->cnt++);
3537
#else
3538
    // free-threading version
3539
    // fast mode uses compare-exchange loop
3540
    // slow mode uses a critical section
3541
    PyObject *returned;
3542
    Py_ssize_t cnt;
3543
3544
    cnt = _Py_atomic_load_ssize_relaxed(&lz->cnt);
3545
    for (;;) {
3546
        if (cnt == PY_SSIZE_T_MAX) {
3547
            Py_BEGIN_CRITICAL_SECTION(lz);
3548
            returned = count_nextlong(lz);
3549
            Py_END_CRITICAL_SECTION();
3550
            return returned;
3551
        }
3552
        if (_Py_atomic_compare_exchange_ssize(&lz->cnt, &cnt, cnt + 1)) {
3553
            return PyLong_FromSsize_t(cnt);
3554
        }
3555
    }
3556
#endif
3557
0
}
3558
3559
static PyObject *
3560
count_repr(PyObject *op)
3561
0
{
3562
0
    countobject *lz = countobject_CAST(op);
3563
0
    if (lz->long_cnt == NULL)
3564
0
        return PyUnicode_FromFormat("%s(%zd)",
3565
0
                                    _PyType_Name(Py_TYPE(lz)), lz->cnt);
3566
3567
0
    if (PyLong_Check(lz->long_step)) {
3568
0
        long step = PyLong_AsLong(lz->long_step);
3569
0
        if (step == -1 && PyErr_Occurred()) {
3570
0
            PyErr_Clear();
3571
0
        }
3572
0
        if (step == 1) {
3573
            /* Don't display step when it is an integer equal to 1 */
3574
0
            return PyUnicode_FromFormat("%s(%R)",
3575
0
                                        _PyType_Name(Py_TYPE(lz)),
3576
0
                                        lz->long_cnt);
3577
0
        }
3578
0
    }
3579
0
    return PyUnicode_FromFormat("%s(%R, %R)",
3580
0
                                _PyType_Name(Py_TYPE(lz)),
3581
0
                                lz->long_cnt, lz->long_step);
3582
0
}
3583
3584
static PyType_Slot count_slots[] = {
3585
    {Py_tp_dealloc, count_dealloc},
3586
    {Py_tp_repr, count_repr},
3587
    {Py_tp_getattro, PyObject_GenericGetAttr},
3588
    {Py_tp_doc, (void *)itertools_count__doc__},
3589
    {Py_tp_traverse, count_traverse},
3590
    {Py_tp_iter, PyObject_SelfIter},
3591
    {Py_tp_iternext, count_next},
3592
    {Py_tp_new, itertools_count},
3593
    {Py_tp_free, PyObject_GC_Del},
3594
    {0, NULL},
3595
};
3596
3597
static PyType_Spec count_spec = {
3598
    .name = "itertools.count",
3599
    .basicsize = sizeof(countobject),
3600
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3601
              Py_TPFLAGS_IMMUTABLETYPE),
3602
    .slots = count_slots,
3603
};
3604
3605
3606
/* repeat object ************************************************************/
3607
3608
typedef struct {
3609
    PyObject_HEAD
3610
    PyObject *element;
3611
    Py_ssize_t cnt;
3612
} repeatobject;
3613
3614
0
#define repeatobject_CAST(op)   ((repeatobject *)(op))
3615
3616
static PyObject *
3617
repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3618
0
{
3619
0
    repeatobject *ro;
3620
0
    PyObject *element;
3621
0
    Py_ssize_t cnt = -1, n_args;
3622
0
    static char *kwargs[] = {"object", "times", NULL};
3623
3624
0
    n_args = PyTuple_GET_SIZE(args);
3625
0
    if (kwds != NULL)
3626
0
        n_args += PyDict_GET_SIZE(kwds);
3627
0
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3628
0
                                     &element, &cnt))
3629
0
        return NULL;
3630
    /* Does user supply times argument? */
3631
0
    if (n_args == 2 && cnt < 0)
3632
0
        cnt = 0;
3633
3634
0
    ro = (repeatobject *)type->tp_alloc(type, 0);
3635
0
    if (ro == NULL)
3636
0
        return NULL;
3637
0
    ro->element = Py_NewRef(element);
3638
0
    ro->cnt = cnt;
3639
0
    return (PyObject *)ro;
3640
0
}
3641
3642
static void
3643
repeat_dealloc(PyObject *op)
3644
0
{
3645
0
    repeatobject *ro = repeatobject_CAST(op);
3646
0
    PyTypeObject *tp = Py_TYPE(ro);
3647
0
    PyObject_GC_UnTrack(ro);
3648
0
    Py_XDECREF(ro->element);
3649
0
    tp->tp_free(ro);
3650
0
    Py_DECREF(tp);
3651
0
}
3652
3653
static int
3654
repeat_traverse(PyObject *op, visitproc visit, void *arg)
3655
0
{
3656
0
    repeatobject *ro = repeatobject_CAST(op);
3657
0
    Py_VISIT(Py_TYPE(ro));
3658
0
    Py_VISIT(ro->element);
3659
0
    return 0;
3660
0
}
3661
3662
static PyObject *
3663
repeat_next(PyObject *op)
3664
0
{
3665
0
    repeatobject *ro = repeatobject_CAST(op);
3666
0
    Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(ro->cnt);
3667
0
    if (cnt == 0) {
3668
0
        return NULL;
3669
0
    }
3670
0
    if (cnt > 0) {
3671
0
        cnt--;
3672
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(ro->cnt, cnt);
3673
0
    }
3674
0
    return Py_NewRef(ro->element);
3675
0
}
3676
3677
static PyObject *
3678
repeat_repr(PyObject *op)
3679
0
{
3680
0
    repeatobject *ro = repeatobject_CAST(op);
3681
0
    if (ro->cnt == -1)
3682
0
        return PyUnicode_FromFormat("%s(%R)",
3683
0
                                    _PyType_Name(Py_TYPE(ro)), ro->element);
3684
0
    else
3685
0
        return PyUnicode_FromFormat("%s(%R, %zd)",
3686
0
                                    _PyType_Name(Py_TYPE(ro)), ro->element,
3687
0
                                    ro->cnt);
3688
0
}
3689
3690
static PyObject *
3691
repeat_len(PyObject *op, PyObject *Py_UNUSED(args))
3692
0
{
3693
0
    repeatobject *ro = repeatobject_CAST(op);
3694
0
    if (ro->cnt == -1) {
3695
0
        PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3696
0
        return NULL;
3697
0
    }
3698
0
    return PyLong_FromSize_t(ro->cnt);
3699
0
}
3700
3701
PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
3702
3703
static PyMethodDef repeat_methods[] = {
3704
    {"__length_hint__", repeat_len, METH_NOARGS, length_hint_doc},
3705
    {NULL,              NULL}           /* sentinel */
3706
};
3707
3708
PyDoc_STRVAR(repeat_doc,
3709
"repeat(object [,times]) -> create an iterator which returns the object\n\
3710
for the specified number of times.  If not specified, returns the object\n\
3711
endlessly.");
3712
3713
static PyType_Slot repeat_slots[] = {
3714
    {Py_tp_dealloc, repeat_dealloc},
3715
    {Py_tp_repr, repeat_repr},
3716
    {Py_tp_getattro, PyObject_GenericGetAttr},
3717
    {Py_tp_doc, (void *)repeat_doc},
3718
    {Py_tp_traverse, repeat_traverse},
3719
    {Py_tp_iter, PyObject_SelfIter},
3720
    {Py_tp_iternext, repeat_next},
3721
    {Py_tp_methods, repeat_methods},
3722
    {Py_tp_new, repeat_new},
3723
    {Py_tp_free, PyObject_GC_Del},
3724
    {0, NULL},
3725
};
3726
3727
static PyType_Spec repeat_spec = {
3728
    .name = "itertools.repeat",
3729
    .basicsize = sizeof(repeatobject),
3730
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3731
              Py_TPFLAGS_IMMUTABLETYPE),
3732
    .slots = repeat_slots,
3733
};
3734
3735
3736
/* ziplongest object *********************************************************/
3737
3738
typedef struct {
3739
    PyObject_HEAD
3740
    Py_ssize_t tuplesize;
3741
    Py_ssize_t numactive;
3742
    PyObject *ittuple;                  /* tuple of iterators */
3743
    PyObject *result;
3744
    PyObject *fillvalue;
3745
} ziplongestobject;
3746
3747
0
#define ziplongestobject_CAST(op)   ((ziplongestobject *)(op))
3748
3749
static PyObject *
3750
zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3751
0
{
3752
0
    ziplongestobject *lz;
3753
0
    Py_ssize_t i;
3754
0
    PyObject *ittuple;  /* tuple of iterators */
3755
0
    PyObject *result;
3756
0
    PyObject *fillvalue = Py_None;
3757
0
    Py_ssize_t tuplesize;
3758
3759
0
    if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
3760
0
        fillvalue = NULL;
3761
0
        if (PyDict_GET_SIZE(kwds) == 1) {
3762
0
            fillvalue = PyDict_GetItemWithError(kwds, &_Py_ID(fillvalue));
3763
0
        }
3764
0
        if (fillvalue == NULL) {
3765
0
            if (!PyErr_Occurred()) {
3766
0
                PyErr_SetString(PyExc_TypeError,
3767
0
                    "zip_longest() got an unexpected keyword argument");
3768
0
            }
3769
0
            return NULL;
3770
0
        }
3771
0
    }
3772
3773
    /* args must be a tuple */
3774
0
    assert(PyTuple_Check(args));
3775
0
    tuplesize = PyTuple_GET_SIZE(args);
3776
3777
    /* obtain iterators */
3778
0
    ittuple = PyTuple_New(tuplesize);
3779
0
    if (ittuple == NULL)
3780
0
        return NULL;
3781
0
    for (i=0; i < tuplesize; i++) {
3782
0
        PyObject *item = PyTuple_GET_ITEM(args, i);
3783
0
        PyObject *it = PyObject_GetIter(item);
3784
0
        if (it == NULL) {
3785
0
            Py_DECREF(ittuple);
3786
0
            return NULL;
3787
0
        }
3788
0
        PyTuple_SET_ITEM(ittuple, i, it);
3789
0
    }
3790
3791
    /* create a result holder */
3792
0
    result = PyTuple_New(tuplesize);
3793
0
    if (result == NULL) {
3794
0
        Py_DECREF(ittuple);
3795
0
        return NULL;
3796
0
    }
3797
0
    for (i=0 ; i < tuplesize ; i++) {
3798
0
        Py_INCREF(Py_None);
3799
0
        PyTuple_SET_ITEM(result, i, Py_None);
3800
0
    }
3801
3802
    /* create ziplongestobject structure */
3803
0
    lz = (ziplongestobject *)type->tp_alloc(type, 0);
3804
0
    if (lz == NULL) {
3805
0
        Py_DECREF(ittuple);
3806
0
        Py_DECREF(result);
3807
0
        return NULL;
3808
0
    }
3809
0
    lz->ittuple = ittuple;
3810
0
    lz->tuplesize = tuplesize;
3811
0
    lz->numactive = tuplesize;
3812
0
    lz->result = result;
3813
0
    lz->fillvalue = Py_NewRef(fillvalue);
3814
0
    return (PyObject *)lz;
3815
0
}
3816
3817
static void
3818
zip_longest_dealloc(PyObject *op)
3819
0
{
3820
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3821
0
    PyTypeObject *tp = Py_TYPE(lz);
3822
0
    PyObject_GC_UnTrack(lz);
3823
0
    Py_XDECREF(lz->ittuple);
3824
0
    Py_XDECREF(lz->result);
3825
0
    Py_XDECREF(lz->fillvalue);
3826
0
    tp->tp_free(lz);
3827
0
    Py_DECREF(tp);
3828
0
}
3829
3830
static int
3831
zip_longest_traverse(PyObject *op, visitproc visit, void *arg)
3832
0
{
3833
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3834
0
    Py_VISIT(Py_TYPE(lz));
3835
0
    Py_VISIT(lz->ittuple);
3836
0
    Py_VISIT(lz->result);
3837
0
    Py_VISIT(lz->fillvalue);
3838
0
    return 0;
3839
0
}
3840
3841
static PyObject *
3842
zip_longest_next(PyObject *op)
3843
0
{
3844
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3845
0
    Py_ssize_t i;
3846
0
    Py_ssize_t tuplesize = lz->tuplesize;
3847
0
    PyObject *result = lz->result;
3848
0
    PyObject *it;
3849
0
    PyObject *item;
3850
0
    PyObject *olditem;
3851
3852
0
    if (tuplesize == 0)
3853
0
        return NULL;
3854
0
    if (lz->numactive == 0)
3855
0
        return NULL;
3856
0
    if (Py_REFCNT(result) == 1) {
3857
0
        Py_INCREF(result);
3858
0
        for (i=0 ; i < tuplesize ; i++) {
3859
0
            it = PyTuple_GET_ITEM(lz->ittuple, i);
3860
0
            if (it == NULL) {
3861
0
                item = Py_NewRef(lz->fillvalue);
3862
0
            } else {
3863
0
                item = PyIter_Next(it);
3864
0
                if (item == NULL) {
3865
0
                    lz->numactive -= 1;
3866
0
                    if (lz->numactive == 0 || PyErr_Occurred()) {
3867
0
                        lz->numactive = 0;
3868
0
                        Py_DECREF(result);
3869
0
                        return NULL;
3870
0
                    } else {
3871
0
                        item = Py_NewRef(lz->fillvalue);
3872
0
                        PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3873
0
                        Py_DECREF(it);
3874
0
                    }
3875
0
                }
3876
0
            }
3877
0
            olditem = PyTuple_GET_ITEM(result, i);
3878
0
            PyTuple_SET_ITEM(result, i, item);
3879
0
            Py_DECREF(olditem);
3880
0
        }
3881
        // bpo-42536: The GC may have untracked this result tuple. Since we're
3882
        // recycling it, make sure it's tracked again:
3883
0
        _PyTuple_Recycle(result);
3884
0
    } else {
3885
0
        result = PyTuple_New(tuplesize);
3886
0
        if (result == NULL)
3887
0
            return NULL;
3888
0
        for (i=0 ; i < tuplesize ; i++) {
3889
0
            it = PyTuple_GET_ITEM(lz->ittuple, i);
3890
0
            if (it == NULL) {
3891
0
                item = Py_NewRef(lz->fillvalue);
3892
0
            } else {
3893
0
                item = PyIter_Next(it);
3894
0
                if (item == NULL) {
3895
0
                    lz->numactive -= 1;
3896
0
                    if (lz->numactive == 0 || PyErr_Occurred()) {
3897
0
                        lz->numactive = 0;
3898
0
                        Py_DECREF(result);
3899
0
                        return NULL;
3900
0
                    } else {
3901
0
                        item = Py_NewRef(lz->fillvalue);
3902
0
                        PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3903
0
                        Py_DECREF(it);
3904
0
                    }
3905
0
                }
3906
0
            }
3907
0
            PyTuple_SET_ITEM(result, i, item);
3908
0
        }
3909
0
    }
3910
0
    return result;
3911
0
}
3912
3913
PyDoc_STRVAR(zip_longest_doc,
3914
"zip_longest(*iterables, fillvalue=None)\n\
3915
--\n\
3916
\n\
3917
Return a zip_longest object whose .__next__() method returns a tuple where\n\
3918
the i-th element comes from the i-th iterable argument.  The .__next__()\n\
3919
method continues until the longest iterable in the argument sequence\n\
3920
is exhausted and then it raises StopIteration.  When the shorter iterables\n\
3921
are exhausted, the fillvalue is substituted in their place.  The fillvalue\n\
3922
defaults to None or can be specified by a keyword argument.\n\
3923
");
3924
3925
static PyType_Slot ziplongest_slots[] = {
3926
    {Py_tp_dealloc, zip_longest_dealloc},
3927
    {Py_tp_getattro, PyObject_GenericGetAttr},
3928
    {Py_tp_doc, (void *)zip_longest_doc},
3929
    {Py_tp_traverse, zip_longest_traverse},
3930
    {Py_tp_iter, PyObject_SelfIter},
3931
    {Py_tp_iternext, zip_longest_next},
3932
    {Py_tp_new, zip_longest_new},
3933
    {Py_tp_free, PyObject_GC_Del},
3934
    {0, NULL},
3935
};
3936
3937
static PyType_Spec ziplongest_spec = {
3938
    .name = "itertools.zip_longest",
3939
    .basicsize = sizeof(ziplongestobject),
3940
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3941
              Py_TPFLAGS_IMMUTABLETYPE),
3942
    .slots = ziplongest_slots,
3943
};
3944
3945
3946
/* module level code ********************************************************/
3947
3948
PyDoc_STRVAR(module_doc,
3949
"Functional tools for creating and using iterators.\n\
3950
\n\
3951
Infinite iterators:\n\
3952
count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
3953
cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
3954
repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
3955
\n\
3956
Iterators terminating on the shortest input sequence:\n\
3957
accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
3958
batched(p, n) --> [p0, p1, ..., p_n-1], [p_n, p_n+1, ..., p_2n-1], ...\n\
3959
chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
3960
chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
3961
compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3962
dropwhile(predicate, seq) --> seq[n], seq[n+1], starting when predicate fails\n\
3963
groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
3964
filterfalse(predicate, seq) --> elements of seq where predicate(elem) is False\n\
3965
islice(seq, [start,] stop [, step]) --> elements from\n\
3966
       seq[start:stop:step]\n\
3967
pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
3968
starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
3969
tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
3970
takewhile(predicate, seq) --> seq[0], seq[1], until predicate fails\n\
3971
zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
3972
\n\
3973
Combinatoric generators:\n\
3974
product(p, q, ... [repeat=1]) --> cartesian product\n\
3975
permutations(p[, r])\n\
3976
combinations(p, r)\n\
3977
combinations_with_replacement(p, r)\n\
3978
");
3979
3980
static int
3981
itertoolsmodule_traverse(PyObject *mod, visitproc visit, void *arg)
3982
5.56k
{
3983
5.56k
    itertools_state *state = get_module_state(mod);
3984
5.56k
    Py_VISIT(state->accumulate_type);
3985
5.56k
    Py_VISIT(state->batched_type);
3986
5.56k
    Py_VISIT(state->chain_type);
3987
5.56k
    Py_VISIT(state->combinations_type);
3988
5.56k
    Py_VISIT(state->compress_type);
3989
5.56k
    Py_VISIT(state->count_type);
3990
5.56k
    Py_VISIT(state->cwr_type);
3991
5.56k
    Py_VISIT(state->cycle_type);
3992
5.56k
    Py_VISIT(state->dropwhile_type);
3993
5.56k
    Py_VISIT(state->filterfalse_type);
3994
5.56k
    Py_VISIT(state->groupby_type);
3995
5.56k
    Py_VISIT(state->_grouper_type);
3996
5.56k
    Py_VISIT(state->islice_type);
3997
5.56k
    Py_VISIT(state->pairwise_type);
3998
5.56k
    Py_VISIT(state->permutations_type);
3999
5.56k
    Py_VISIT(state->product_type);
4000
5.56k
    Py_VISIT(state->repeat_type);
4001
5.56k
    Py_VISIT(state->starmap_type);
4002
5.56k
    Py_VISIT(state->takewhile_type);
4003
5.56k
    Py_VISIT(state->tee_type);
4004
5.56k
    Py_VISIT(state->teedataobject_type);
4005
5.56k
    Py_VISIT(state->ziplongest_type);
4006
5.56k
    return 0;
4007
5.56k
}
4008
4009
static int
4010
itertoolsmodule_clear(PyObject *mod)
4011
0
{
4012
0
    itertools_state *state = get_module_state(mod);
4013
0
    Py_CLEAR(state->accumulate_type);
4014
0
    Py_CLEAR(state->batched_type);
4015
0
    Py_CLEAR(state->chain_type);
4016
0
    Py_CLEAR(state->combinations_type);
4017
0
    Py_CLEAR(state->compress_type);
4018
0
    Py_CLEAR(state->count_type);
4019
0
    Py_CLEAR(state->cwr_type);
4020
0
    Py_CLEAR(state->cycle_type);
4021
0
    Py_CLEAR(state->dropwhile_type);
4022
0
    Py_CLEAR(state->filterfalse_type);
4023
0
    Py_CLEAR(state->groupby_type);
4024
0
    Py_CLEAR(state->_grouper_type);
4025
0
    Py_CLEAR(state->islice_type);
4026
0
    Py_CLEAR(state->pairwise_type);
4027
0
    Py_CLEAR(state->permutations_type);
4028
0
    Py_CLEAR(state->product_type);
4029
0
    Py_CLEAR(state->repeat_type);
4030
0
    Py_CLEAR(state->starmap_type);
4031
0
    Py_CLEAR(state->takewhile_type);
4032
0
    Py_CLEAR(state->tee_type);
4033
0
    Py_CLEAR(state->teedataobject_type);
4034
0
    Py_CLEAR(state->ziplongest_type);
4035
0
    return 0;
4036
0
}
4037
4038
static void
4039
itertoolsmodule_free(void *mod)
4040
0
{
4041
0
    (void)itertoolsmodule_clear((PyObject *)mod);
4042
0
}
4043
4044
264
#define ADD_TYPE(module, type, spec)                                     \
4045
264
do {                                                                     \
4046
264
    type = (PyTypeObject *)PyType_FromModuleAndSpec(module, spec, NULL); \
4047
264
    if (type == NULL) {                                                  \
4048
0
        return -1;                                                       \
4049
0
    }                                                                    \
4050
264
    if (PyModule_AddType(module, type) < 0) {                            \
4051
0
        return -1;                                                       \
4052
0
    }                                                                    \
4053
264
} while (0)
4054
4055
static int
4056
itertoolsmodule_exec(PyObject *mod)
4057
12
{
4058
12
    itertools_state *state = get_module_state(mod);
4059
12
    ADD_TYPE(mod, state->accumulate_type, &accumulate_spec);
4060
12
    ADD_TYPE(mod, state->batched_type, &batched_spec);
4061
12
    ADD_TYPE(mod, state->chain_type, &chain_spec);
4062
12
    ADD_TYPE(mod, state->combinations_type, &combinations_spec);
4063
12
    ADD_TYPE(mod, state->compress_type, &compress_spec);
4064
12
    ADD_TYPE(mod, state->count_type, &count_spec);
4065
12
    ADD_TYPE(mod, state->cwr_type, &cwr_spec);
4066
12
    ADD_TYPE(mod, state->cycle_type, &cycle_spec);
4067
12
    ADD_TYPE(mod, state->dropwhile_type, &dropwhile_spec);
4068
12
    ADD_TYPE(mod, state->filterfalse_type, &filterfalse_spec);
4069
12
    ADD_TYPE(mod, state->groupby_type, &groupby_spec);
4070
12
    ADD_TYPE(mod, state->_grouper_type, &_grouper_spec);
4071
12
    ADD_TYPE(mod, state->islice_type, &islice_spec);
4072
12
    ADD_TYPE(mod, state->pairwise_type, &pairwise_spec);
4073
12
    ADD_TYPE(mod, state->permutations_type, &permutations_spec);
4074
12
    ADD_TYPE(mod, state->product_type, &product_spec);
4075
12
    ADD_TYPE(mod, state->repeat_type, &repeat_spec);
4076
12
    ADD_TYPE(mod, state->starmap_type, &starmap_spec);
4077
12
    ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
4078
12
    ADD_TYPE(mod, state->tee_type, &tee_spec);
4079
12
    ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
4080
12
    ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
4081
4082
12
    Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
4083
12
    return 0;
4084
12
}
4085
4086
static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4087
    {Py_mod_exec, itertoolsmodule_exec},
4088
    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
4089
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
4090
    {0, NULL}
4091
};
4092
4093
static PyMethodDef module_methods[] = {
4094
    ITERTOOLS_TEE_METHODDEF
4095
    {NULL, NULL} /* sentinel */
4096
};
4097
4098
4099
static struct PyModuleDef itertoolsmodule = {
4100
    .m_base = PyModuleDef_HEAD_INIT,
4101
    .m_name = "itertools",
4102
    .m_doc = module_doc,
4103
    .m_size = sizeof(itertools_state),
4104
    .m_methods = module_methods,
4105
    .m_slots = itertoolsmodule_slots,
4106
    .m_traverse = itertoolsmodule_traverse,
4107
    .m_clear = itertoolsmodule_clear,
4108
    .m_free = itertoolsmodule_free,
4109
};
4110
4111
PyMODINIT_FUNC
4112
PyInit_itertools(void)
4113
12
{
4114
12
    return PyModuleDef_Init(&itertoolsmodule);
4115
12
}