Coverage Report

Created: 2026-02-09 07:07

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