Coverage Report

Created: 2025-09-05 07:10

/src/cpython/Modules/itertoolsmodule.c
Line
Count
Source (jump to first uncovered line)
1
#include "Python.h"
2
#include "pycore_call.h"              // _PyObject_CallNoArgs()
3
#include "pycore_ceval.h"             // _PyEval_GetBuiltin()
4
#include "pycore_critical_section.h"  // Py_BEGIN_CRITICAL_SECTION()
5
#include "pycore_long.h"              // _PyLong_GetZero()
6
#include "pycore_moduleobject.h"      // _PyModule_GetState()
7
#include "pycore_typeobject.h"        // _PyType_GetModuleState()
8
#include "pycore_object.h"            // _PyObject_GC_TRACK()
9
#include "pycore_tuple.h"             // _PyTuple_ITEMS()
10
11
#include <stddef.h>                   // offsetof()
12
13
/* Itertools module written and maintained
14
   by Raymond D. Hettinger <python@rcn.com>
15
*/
16
17
typedef struct {
18
    PyTypeObject *accumulate_type;
19
    PyTypeObject *batched_type;
20
    PyTypeObject *chain_type;
21
    PyTypeObject *combinations_type;
22
    PyTypeObject *compress_type;
23
    PyTypeObject *count_type;
24
    PyTypeObject *cwr_type;
25
    PyTypeObject *cycle_type;
26
    PyTypeObject *dropwhile_type;
27
    PyTypeObject *filterfalse_type;
28
    PyTypeObject *groupby_type;
29
    PyTypeObject *_grouper_type;
30
    PyTypeObject *islice_type;
31
    PyTypeObject *pairwise_type;
32
    PyTypeObject *permutations_type;
33
    PyTypeObject *product_type;
34
    PyTypeObject *repeat_type;
35
    PyTypeObject *starmap_type;
36
    PyTypeObject *takewhile_type;
37
    PyTypeObject *tee_type;
38
    PyTypeObject *teedataobject_type;
39
    PyTypeObject *ziplongest_type;
40
} itertools_state;
41
42
static inline itertools_state *
43
get_module_state(PyObject *mod)
44
5.27k
{
45
5.27k
    void *state = _PyModule_GetState(mod);
46
5.27k
    assert(state != NULL);
47
5.27k
    return (itertools_state *)state;
48
5.27k
}
49
50
static inline itertools_state *
51
get_module_state_by_cls(PyTypeObject *cls)
52
0
{
53
0
    void *state = _PyType_GetModuleState(cls);
54
0
    assert(state != NULL);
55
0
    return (itertools_state *)state;
56
0
}
57
58
static struct PyModuleDef itertoolsmodule;
59
60
static inline itertools_state *
61
find_state_by_type(PyTypeObject *tp)
62
0
{
63
0
    PyObject *mod = PyType_GetModuleByDef(tp, &itertoolsmodule);
64
0
    assert(mod != NULL);
65
0
    return get_module_state(mod);
66
0
}
67
68
/*[clinic input]
69
module itertools
70
class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
71
class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
72
class itertools.teedataobject "teedataobject *" "clinic_state()->teedataobject_type"
73
class itertools._tee "teeobject *" "clinic_state()->tee_type"
74
class itertools.batched "batchedobject *" "clinic_state()->batched_type"
75
class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
76
class itertools.dropwhile "dropwhileobject *" "clinic_state()->dropwhile_type"
77
class itertools.takewhile "takewhileobject *" "clinic_state()->takewhile_type"
78
class itertools.starmap "starmapobject *" "clinic_state()->starmap_type"
79
class itertools.chain "chainobject *" "clinic_state()->chain_type"
80
class itertools.combinations "combinationsobject *" "clinic_state()->combinations_type"
81
class itertools.combinations_with_replacement "cwr_object *" "clinic_state()->cwr_type"
82
class itertools.permutations "permutationsobject *" "clinic_state()->permutations_type"
83
class itertools.accumulate "accumulateobject *" "clinic_state()->accumulate_type"
84
class itertools.compress "compressobject *" "clinic_state()->compress_type"
85
class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_type"
86
class itertools.count "countobject *" "clinic_state()->count_type"
87
class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
88
[clinic start generated code]*/
89
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=aa48fe4de9d4080f]*/
90
91
0
#define clinic_state() (find_state_by_type(type))
92
0
#define clinic_state_by_cls() (get_module_state_by_cls(base_tp))
93
#include "clinic/itertoolsmodule.c.h"
94
#undef clinic_state_by_cls
95
#undef clinic_state
96
97
98
/* batched object ************************************************************/
99
100
typedef struct {
101
    PyObject_HEAD
102
    PyObject *it;
103
    Py_ssize_t batch_size;
104
    bool strict;
105
} batchedobject;
106
107
0
#define batchedobject_CAST(op)  ((batchedobject *)(op))
108
109
/*[clinic input]
110
@classmethod
111
itertools.batched.__new__ as batched_new
112
    iterable: object
113
    n: Py_ssize_t
114
    *
115
    strict: bool = False
116
117
Batch data into tuples of length n. The last batch may be shorter than n.
118
119
Loops over the input iterable and accumulates data into tuples
120
up to size n.  The input is consumed lazily, just enough to
121
fill a batch.  The result is yielded as soon as a batch is full
122
or when the input iterable is exhausted.
123
124
    >>> for batch in batched('ABCDEFG', 3):
125
    ...     print(batch)
126
    ...
127
    ('A', 'B', 'C')
128
    ('D', 'E', 'F')
129
    ('G',)
130
131
If "strict" is True, raises a ValueError if the final batch is shorter
132
than n.
133
134
[clinic start generated code]*/
135
136
static PyObject *
137
batched_new_impl(PyTypeObject *type, PyObject *iterable, Py_ssize_t n,
138
                 int strict)
139
/*[clinic end generated code: output=c6de11b061529d3e input=7814b47e222f5467]*/
140
0
{
141
0
    PyObject *it;
142
0
    batchedobject *bo;
143
144
0
    if (n < 1) {
145
        /* We could define the n==0 case to return an empty iterator
146
           but that is at odds with the idea that batching should
147
           never throw-away input data.
148
        */
149
0
        PyErr_SetString(PyExc_ValueError, "n must be at least one");
150
0
        return NULL;
151
0
    }
152
0
    it = PyObject_GetIter(iterable);
153
0
    if (it == NULL) {
154
0
        return NULL;
155
0
    }
156
157
    /* create batchedobject structure */
158
0
    bo = (batchedobject *)type->tp_alloc(type, 0);
159
0
    if (bo == NULL) {
160
0
        Py_DECREF(it);
161
0
        return NULL;
162
0
    }
163
0
    bo->batch_size = n;
164
0
    bo->it = it;
165
0
    bo->strict = (bool) strict;
166
0
    return (PyObject *)bo;
167
0
}
168
169
static void
170
batched_dealloc(PyObject *op)
171
0
{
172
0
    batchedobject *bo = batchedobject_CAST(op);
173
0
    PyTypeObject *tp = Py_TYPE(bo);
174
0
    PyObject_GC_UnTrack(bo);
175
0
    Py_XDECREF(bo->it);
176
0
    tp->tp_free(bo);
177
0
    Py_DECREF(tp);
178
0
}
179
180
static int
181
batched_traverse(PyObject *op, visitproc visit, void *arg)
182
0
{
183
0
    batchedobject *bo = batchedobject_CAST(op);
184
0
    Py_VISIT(Py_TYPE(bo));
185
0
    Py_VISIT(bo->it);
186
0
    return 0;
187
0
}
188
189
static PyObject *
190
batched_next(PyObject *op)
191
0
{
192
0
    batchedobject *bo = batchedobject_CAST(op);
193
0
    Py_ssize_t i;
194
0
    Py_ssize_t n = FT_ATOMIC_LOAD_SSIZE_RELAXED(bo->batch_size);
195
0
    PyObject *it = bo->it;
196
0
    PyObject *item;
197
0
    PyObject *result;
198
199
0
    if (n < 0) {
200
0
        return NULL;
201
0
    }
202
0
    result = PyTuple_New(n);
203
0
    if (result == NULL) {
204
0
        return NULL;
205
0
    }
206
0
    iternextfunc iternext = *Py_TYPE(it)->tp_iternext;
207
0
    PyObject **items = _PyTuple_ITEMS(result);
208
0
    for (i=0 ; i < n ; i++) {
209
0
        item = iternext(it);
210
0
        if (item == NULL) {
211
0
            goto null_item;
212
0
        }
213
0
        items[i] = item;
214
0
    }
215
0
    return result;
216
217
0
 null_item:
218
0
    if (PyErr_Occurred()) {
219
0
        if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
220
            /* Input raised an exception other than StopIteration */
221
0
            FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
222
0
#ifndef Py_GIL_DISABLED
223
0
            Py_CLEAR(bo->it);
224
0
#endif
225
0
            Py_DECREF(result);
226
0
            return NULL;
227
0
        }
228
0
        PyErr_Clear();
229
0
    }
230
0
    if (i == 0) {
231
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
232
0
#ifndef Py_GIL_DISABLED
233
0
        Py_CLEAR(bo->it);
234
0
#endif
235
0
        Py_DECREF(result);
236
0
        return NULL;
237
0
    }
238
0
    if (bo->strict) {
239
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(bo->batch_size, -1);
240
0
#ifndef Py_GIL_DISABLED
241
0
        Py_CLEAR(bo->it);
242
0
#endif
243
0
        Py_DECREF(result);
244
0
        PyErr_SetString(PyExc_ValueError, "batched(): incomplete batch");
245
0
        return NULL;
246
0
    }
247
0
    _PyTuple_Resize(&result, i);
248
0
    return result;
249
0
}
250
251
static PyType_Slot batched_slots[] = {
252
    {Py_tp_dealloc, batched_dealloc},
253
    {Py_tp_getattro, PyObject_GenericGetAttr},
254
    {Py_tp_doc, (void *)batched_new__doc__},
255
    {Py_tp_traverse, batched_traverse},
256
    {Py_tp_iter, PyObject_SelfIter},
257
    {Py_tp_iternext, batched_next},
258
    {Py_tp_alloc, PyType_GenericAlloc},
259
    {Py_tp_new, batched_new},
260
    {Py_tp_free, PyObject_GC_Del},
261
    {0, NULL},
262
};
263
264
static PyType_Spec batched_spec = {
265
    .name = "itertools.batched",
266
    .basicsize = sizeof(batchedobject),
267
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
268
              Py_TPFLAGS_IMMUTABLETYPE),
269
    .slots = batched_slots,
270
};
271
272
273
/* pairwise object ***********************************************************/
274
275
typedef struct {
276
    PyObject_HEAD
277
    PyObject *it;
278
    PyObject *old;
279
    PyObject *result;
280
} pairwiseobject;
281
282
0
#define pairwiseobject_CAST(op) ((pairwiseobject *)(op))
283
284
/*[clinic input]
285
@classmethod
286
itertools.pairwise.__new__ as pairwise_new
287
    iterable: object
288
    /
289
Return an iterator of overlapping pairs taken from the input iterator.
290
291
    s -> (s0,s1), (s1,s2), (s2, s3), ...
292
293
[clinic start generated code]*/
294
295
static PyObject *
296
pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
297
/*[clinic end generated code: output=9f0267062d384456 input=6e7c3cddb431a8d6]*/
298
0
{
299
0
    PyObject *it;
300
0
    pairwiseobject *po;
301
302
0
    it = PyObject_GetIter(iterable);
303
0
    if (it == NULL) {
304
0
        return NULL;
305
0
    }
306
0
    po = (pairwiseobject *)type->tp_alloc(type, 0);
307
0
    if (po == NULL) {
308
0
        Py_DECREF(it);
309
0
        return NULL;
310
0
    }
311
0
    po->it = it;
312
0
    po->old = NULL;
313
0
    po->result = PyTuple_Pack(2, Py_None, Py_None);
314
0
    if (po->result == NULL) {
315
0
        Py_DECREF(po);
316
0
        return NULL;
317
0
    }
318
0
    return (PyObject *)po;
319
0
}
320
321
static void
322
pairwise_dealloc(PyObject *op)
323
0
{
324
0
    pairwiseobject *po = pairwiseobject_CAST(op);
325
0
    PyTypeObject *tp = Py_TYPE(po);
326
0
    PyObject_GC_UnTrack(po);
327
0
    Py_XDECREF(po->it);
328
0
    Py_XDECREF(po->old);
329
0
    Py_XDECREF(po->result);
330
0
    tp->tp_free(po);
331
0
    Py_DECREF(tp);
332
0
}
333
334
static int
335
pairwise_traverse(PyObject *op, visitproc visit, void *arg)
336
0
{
337
0
    pairwiseobject *po = pairwiseobject_CAST(op);
338
0
    Py_VISIT(Py_TYPE(po));
339
0
    Py_VISIT(po->it);
340
0
    Py_VISIT(po->old);
341
0
    Py_VISIT(po->result);
342
0
    return 0;
343
0
}
344
345
static PyObject *
346
pairwise_next(PyObject *op)
347
0
{
348
0
    pairwiseobject *po = pairwiseobject_CAST(op);
349
0
    PyObject *it = po->it;
350
0
    PyObject *old = po->old;
351
0
    PyObject *new, *result;
352
353
0
    if (it == NULL) {
354
0
        return NULL;
355
0
    }
356
0
    if (old == NULL) {
357
0
        old = (*Py_TYPE(it)->tp_iternext)(it);
358
0
        Py_XSETREF(po->old, old);
359
0
        if (old == NULL) {
360
0
            Py_CLEAR(po->it);
361
0
            return NULL;
362
0
        }
363
0
        it = po->it;
364
0
        if (it == NULL) {
365
0
            Py_CLEAR(po->old);
366
0
            return NULL;
367
0
        }
368
0
    }
369
0
    Py_INCREF(old);
370
0
    new = (*Py_TYPE(it)->tp_iternext)(it);
371
0
    if (new == NULL) {
372
0
        Py_CLEAR(po->it);
373
0
        Py_CLEAR(po->old);
374
0
        Py_DECREF(old);
375
0
        return NULL;
376
0
    }
377
378
0
    result = po->result;
379
0
    if (Py_REFCNT(result) == 1) {
380
0
        Py_INCREF(result);
381
0
        PyObject *last_old = PyTuple_GET_ITEM(result, 0);
382
0
        PyObject *last_new = PyTuple_GET_ITEM(result, 1);
383
0
        PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
384
0
        PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
385
0
        Py_DECREF(last_old);
386
0
        Py_DECREF(last_new);
387
        // bpo-42536: The GC may have untracked this result tuple. Since we're
388
        // recycling it, make sure it's tracked again:
389
0
        _PyTuple_Recycle(result);
390
0
    }
391
0
    else {
392
0
        result = PyTuple_New(2);
393
0
        if (result != NULL) {
394
0
            PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
395
0
            PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
396
0
        }
397
0
    }
398
399
0
    Py_XSETREF(po->old, new);
400
0
    Py_DECREF(old);
401
0
    return result;
402
0
}
403
404
static PyType_Slot pairwise_slots[] = {
405
    {Py_tp_dealloc, pairwise_dealloc},
406
    {Py_tp_getattro, PyObject_GenericGetAttr},
407
    {Py_tp_doc, (void *)pairwise_new__doc__},
408
    {Py_tp_traverse, pairwise_traverse},
409
    {Py_tp_iter, PyObject_SelfIter},
410
    {Py_tp_iternext, pairwise_next},
411
    {Py_tp_alloc, PyType_GenericAlloc},
412
    {Py_tp_new, pairwise_new},
413
    {Py_tp_free, PyObject_GC_Del},
414
    {0, NULL},
415
};
416
417
static PyType_Spec pairwise_spec = {
418
    .name = "itertools.pairwise",
419
    .basicsize = sizeof(pairwiseobject),
420
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
421
              Py_TPFLAGS_IMMUTABLETYPE),
422
    .slots = pairwise_slots,
423
};
424
425
426
/* groupby object ************************************************************/
427
428
typedef struct {
429
    PyObject_HEAD
430
    PyObject *it;
431
    PyObject *keyfunc;
432
    PyObject *tgtkey;
433
    PyObject *currkey;
434
    PyObject *currvalue;
435
    const void *currgrouper;  /* borrowed reference */
436
    itertools_state *state;
437
} groupbyobject;
438
439
0
#define groupbyobject_CAST(op)  ((groupbyobject *)(op))
440
441
static PyObject *_grouper_create(groupbyobject *, PyObject *);
442
443
/*[clinic input]
444
@classmethod
445
itertools.groupby.__new__
446
447
    iterable as it: object
448
        Elements to divide into groups according to the key function.
449
    key as keyfunc: object = None
450
        A function for computing the group category for each element.
451
        If the key function is not specified or is None, the element itself
452
        is used for grouping.
453
454
make an iterator that returns consecutive keys and groups from the iterable
455
[clinic start generated code]*/
456
457
static PyObject *
458
itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
459
/*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
460
0
{
461
0
    groupbyobject *gbo;
462
463
0
    gbo = (groupbyobject *)type->tp_alloc(type, 0);
464
0
    if (gbo == NULL)
465
0
        return NULL;
466
0
    gbo->tgtkey = NULL;
467
0
    gbo->currkey = NULL;
468
0
    gbo->currvalue = NULL;
469
0
    gbo->keyfunc = Py_NewRef(keyfunc);
470
0
    gbo->it = PyObject_GetIter(it);
471
0
    if (gbo->it == NULL) {
472
0
        Py_DECREF(gbo);
473
0
        return NULL;
474
0
    }
475
0
    gbo->state = find_state_by_type(type);
476
0
    return (PyObject *)gbo;
477
0
}
478
479
static void
480
groupby_dealloc(PyObject *op)
481
0
{
482
0
    groupbyobject *gbo = groupbyobject_CAST(op);
483
0
    PyTypeObject *tp = Py_TYPE(gbo);
484
0
    PyObject_GC_UnTrack(gbo);
485
0
    Py_XDECREF(gbo->it);
486
0
    Py_XDECREF(gbo->keyfunc);
487
0
    Py_XDECREF(gbo->tgtkey);
488
0
    Py_XDECREF(gbo->currkey);
489
0
    Py_XDECREF(gbo->currvalue);
490
0
    tp->tp_free(gbo);
491
0
    Py_DECREF(tp);
492
0
}
493
494
static int
495
groupby_traverse(PyObject *op, visitproc visit, void *arg)
496
0
{
497
0
    groupbyobject *gbo = groupbyobject_CAST(op);
498
0
    Py_VISIT(Py_TYPE(gbo));
499
0
    Py_VISIT(gbo->it);
500
0
    Py_VISIT(gbo->keyfunc);
501
0
    Py_VISIT(gbo->tgtkey);
502
0
    Py_VISIT(gbo->currkey);
503
0
    Py_VISIT(gbo->currvalue);
504
0
    return 0;
505
0
}
506
507
Py_LOCAL_INLINE(int)
508
groupby_step(groupbyobject *gbo)
509
0
{
510
0
    PyObject *newvalue, *newkey, *oldvalue;
511
512
0
    newvalue = PyIter_Next(gbo->it);
513
0
    if (newvalue == NULL)
514
0
        return -1;
515
516
0
    if (gbo->keyfunc == Py_None) {
517
0
        newkey = Py_NewRef(newvalue);
518
0
    } else {
519
0
        newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
520
0
        if (newkey == NULL) {
521
0
            Py_DECREF(newvalue);
522
0
            return -1;
523
0
        }
524
0
    }
525
526
0
    oldvalue = gbo->currvalue;
527
0
    gbo->currvalue = newvalue;
528
0
    Py_XSETREF(gbo->currkey, newkey);
529
0
    Py_XDECREF(oldvalue);
530
0
    return 0;
531
0
}
532
533
static PyObject *
534
groupby_next(PyObject *op)
535
0
{
536
0
    PyObject *r, *grouper;
537
0
    groupbyobject *gbo = groupbyobject_CAST(op);
538
539
0
    gbo->currgrouper = NULL;
540
    /* skip to next iteration group */
541
0
    for (;;) {
542
0
        if (gbo->currkey == NULL)
543
0
            /* pass */;
544
0
        else if (gbo->tgtkey == NULL)
545
0
            break;
546
0
        else {
547
0
            int rcmp;
548
549
0
            rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
550
0
            if (rcmp == -1)
551
0
                return NULL;
552
0
            else if (rcmp == 0)
553
0
                break;
554
0
        }
555
556
0
        if (groupby_step(gbo) < 0)
557
0
            return NULL;
558
0
    }
559
0
    Py_INCREF(gbo->currkey);
560
0
    Py_XSETREF(gbo->tgtkey, gbo->currkey);
561
562
0
    grouper = _grouper_create(gbo, gbo->tgtkey);
563
0
    if (grouper == NULL)
564
0
        return NULL;
565
566
0
    r = PyTuple_Pack(2, gbo->currkey, grouper);
567
0
    Py_DECREF(grouper);
568
0
    return r;
569
0
}
570
571
static PyType_Slot groupby_slots[] = {
572
    {Py_tp_dealloc, groupby_dealloc},
573
    {Py_tp_getattro, PyObject_GenericGetAttr},
574
    {Py_tp_doc, (void *)itertools_groupby__doc__},
575
    {Py_tp_traverse, groupby_traverse},
576
    {Py_tp_iter, PyObject_SelfIter},
577
    {Py_tp_iternext, groupby_next},
578
    {Py_tp_new, itertools_groupby},
579
    {Py_tp_free, PyObject_GC_Del},
580
    {0, NULL},
581
};
582
583
static PyType_Spec groupby_spec = {
584
    .name = "itertools.groupby",
585
    .basicsize= sizeof(groupbyobject),
586
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
587
              Py_TPFLAGS_IMMUTABLETYPE),
588
    .slots = groupby_slots,
589
};
590
591
/* _grouper object (internal) ************************************************/
592
593
typedef struct {
594
    PyObject_HEAD
595
    PyObject *parent;
596
    PyObject *tgtkey;
597
} _grouperobject;
598
599
0
#define _grouperobject_CAST(op) ((_grouperobject *)(op))
600
601
/*[clinic input]
602
@classmethod
603
itertools._grouper.__new__
604
605
    parent: object(subclass_of='clinic_state_by_cls()->groupby_type')
606
    tgtkey: object
607
    /
608
[clinic start generated code]*/
609
610
static PyObject *
611
itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
612
                        PyObject *tgtkey)
613
/*[clinic end generated code: output=462efb1cdebb5914 input=afe05eb477118f12]*/
614
0
{
615
0
    return _grouper_create(groupbyobject_CAST(parent), tgtkey);
616
0
}
617
618
static PyObject *
619
_grouper_create(groupbyobject *parent, PyObject *tgtkey)
620
0
{
621
0
    itertools_state *state = parent->state;
622
0
    _grouperobject *igo = PyObject_GC_New(_grouperobject, state->_grouper_type);
623
0
    if (igo == NULL)
624
0
        return NULL;
625
0
    igo->parent = Py_NewRef(parent);
626
0
    igo->tgtkey = Py_NewRef(tgtkey);
627
0
    parent->currgrouper = igo;  /* borrowed reference */
628
629
0
    PyObject_GC_Track(igo);
630
0
    return (PyObject *)igo;
631
0
}
632
633
static void
634
_grouper_dealloc(PyObject *op)
635
0
{
636
0
    _grouperobject *igo = _grouperobject_CAST(op);
637
0
    PyTypeObject *tp = Py_TYPE(igo);
638
0
    PyObject_GC_UnTrack(igo);
639
0
    Py_DECREF(igo->parent);
640
0
    Py_DECREF(igo->tgtkey);
641
0
    PyObject_GC_Del(igo);
642
0
    Py_DECREF(tp);
643
0
}
644
645
static int
646
_grouper_traverse(PyObject *op, visitproc visit, void *arg)
647
0
{
648
0
    _grouperobject *igo = _grouperobject_CAST(op);
649
0
    Py_VISIT(Py_TYPE(igo));
650
0
    Py_VISIT(igo->parent);
651
0
    Py_VISIT(igo->tgtkey);
652
0
    return 0;
653
0
}
654
655
static PyObject *
656
_grouper_next(PyObject *op)
657
0
{
658
0
    _grouperobject *igo = _grouperobject_CAST(op);
659
0
    groupbyobject *gbo = groupbyobject_CAST(igo->parent);
660
0
    PyObject *r;
661
0
    int rcmp;
662
663
0
    if (gbo->currgrouper != igo)
664
0
        return NULL;
665
0
    if (gbo->currvalue == NULL) {
666
0
        if (groupby_step(gbo) < 0)
667
0
            return NULL;
668
0
    }
669
670
0
    assert(gbo->currkey != NULL);
671
0
    rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
672
0
    if (rcmp <= 0)
673
        /* got any error or current group is end */
674
0
        return NULL;
675
676
0
    r = gbo->currvalue;
677
0
    gbo->currvalue = NULL;
678
0
    Py_CLEAR(gbo->currkey);
679
680
0
    return r;
681
0
}
682
683
static PyType_Slot _grouper_slots[] = {
684
    {Py_tp_dealloc, _grouper_dealloc},
685
    {Py_tp_getattro, PyObject_GenericGetAttr},
686
    {Py_tp_traverse, _grouper_traverse},
687
    {Py_tp_iter, PyObject_SelfIter},
688
    {Py_tp_iternext, _grouper_next},
689
    {Py_tp_new, itertools__grouper},
690
    {Py_tp_free, PyObject_GC_Del},
691
    {0, NULL},
692
};
693
694
static PyType_Spec _grouper_spec = {
695
    .name = "itertools._grouper",
696
    .basicsize = sizeof(_grouperobject),
697
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
698
              Py_TPFLAGS_IMMUTABLETYPE),
699
    .slots = _grouper_slots,
700
};
701
702
703
/* tee object and with supporting function and objects ***********************/
704
705
/* The teedataobject pre-allocates space for LINKCELLS number of objects.
706
   To help the object fit neatly inside cache lines (space for 16 to 32
707
   pointers), the value should be a multiple of 16 minus  space for
708
   the other structure members including PyHEAD overhead.  The larger the
709
   value, the less memory overhead per object and the less time spent
710
   allocating/deallocating new links.  The smaller the number, the less
711
   wasted space and the more rapid freeing of older data.
712
*/
713
0
#define LINKCELLS 57
714
715
typedef struct {
716
    PyObject_HEAD
717
    PyObject *it;
718
    int numread;                /* 0 <= numread <= LINKCELLS */
719
    int running;
720
    PyObject *nextlink;
721
    PyObject *(values[LINKCELLS]);
722
} teedataobject;
723
724
0
#define teedataobject_CAST(op)  ((teedataobject *)(op))
725
726
typedef struct {
727
    PyObject_HEAD
728
    teedataobject *dataobj;
729
    int index;                  /* 0 <= index <= LINKCELLS */
730
    PyObject *weakreflist;
731
    itertools_state *state;
732
} teeobject;
733
734
0
#define teeobject_CAST(op)  ((teeobject *)(op))
735
736
static PyObject *
737
teedataobject_newinternal(itertools_state *state, PyObject *it)
738
0
{
739
0
    teedataobject *tdo;
740
741
0
    tdo = PyObject_GC_New(teedataobject, state->teedataobject_type);
742
0
    if (tdo == NULL)
743
0
        return NULL;
744
745
0
    tdo->running = 0;
746
0
    tdo->numread = 0;
747
0
    tdo->nextlink = NULL;
748
0
    tdo->it = Py_NewRef(it);
749
0
    PyObject_GC_Track(tdo);
750
0
    return (PyObject *)tdo;
751
0
}
752
753
static PyObject *
754
teedataobject_jumplink(itertools_state *state, teedataobject *tdo)
755
0
{
756
0
    if (tdo->nextlink == NULL)
757
0
        tdo->nextlink = teedataobject_newinternal(state, tdo->it);
758
0
    return Py_XNewRef(tdo->nextlink);
759
0
}
760
761
static PyObject *
762
teedataobject_getitem(teedataobject *tdo, int i)
763
0
{
764
0
    PyObject *value;
765
766
0
    assert(i < LINKCELLS);
767
0
    if (i < tdo->numread)
768
0
        value = tdo->values[i];
769
0
    else {
770
        /* this is the lead iterator, so fetch more data */
771
0
        assert(i == tdo->numread);
772
0
        if (tdo->running) {
773
0
            PyErr_SetString(PyExc_RuntimeError,
774
0
                            "cannot re-enter the tee iterator");
775
0
            return NULL;
776
0
        }
777
0
        tdo->running = 1;
778
0
        value = PyIter_Next(tdo->it);
779
0
        tdo->running = 0;
780
0
        if (value == NULL)
781
0
            return NULL;
782
0
        tdo->numread++;
783
0
        tdo->values[i] = value;
784
0
    }
785
0
    return Py_NewRef(value);
786
0
}
787
788
static int
789
teedataobject_traverse(PyObject *op, visitproc visit, void * arg)
790
0
{
791
0
    int i;
792
0
    teedataobject *tdo = teedataobject_CAST(op);
793
794
0
    Py_VISIT(Py_TYPE(tdo));
795
0
    Py_VISIT(tdo->it);
796
0
    for (i = 0; i < tdo->numread; i++)
797
0
        Py_VISIT(tdo->values[i]);
798
0
    Py_VISIT(tdo->nextlink);
799
0
    return 0;
800
0
}
801
802
static void
803
teedataobject_safe_decref(PyObject *obj)
804
0
{
805
0
    while (obj && Py_REFCNT(obj) == 1) {
806
0
        teedataobject *tmp = teedataobject_CAST(obj);
807
0
        PyObject *nextlink = tmp->nextlink;
808
0
        tmp->nextlink = NULL;
809
0
        Py_SETREF(obj, nextlink);
810
0
    }
811
0
    Py_XDECREF(obj);
812
0
}
813
814
static int
815
teedataobject_clear(PyObject *op)
816
0
{
817
0
    int i;
818
0
    PyObject *tmp;
819
0
    teedataobject *tdo = teedataobject_CAST(op);
820
821
0
    Py_CLEAR(tdo->it);
822
0
    for (i=0 ; i<tdo->numread ; i++)
823
0
        Py_CLEAR(tdo->values[i]);
824
0
    tmp = tdo->nextlink;
825
0
    tdo->nextlink = NULL;
826
0
    teedataobject_safe_decref(tmp);
827
0
    return 0;
828
0
}
829
830
static void
831
teedataobject_dealloc(PyObject *op)
832
0
{
833
0
    PyTypeObject *tp = Py_TYPE(op);
834
0
    PyObject_GC_UnTrack(op);
835
0
    (void)teedataobject_clear(op);
836
0
    PyObject_GC_Del(op);
837
0
    Py_DECREF(tp);
838
0
}
839
840
/*[clinic input]
841
@classmethod
842
itertools.teedataobject.__new__
843
    iterable as it: object
844
    values: object(subclass_of='&PyList_Type')
845
    next: object
846
    /
847
Data container common to multiple tee objects.
848
[clinic start generated code]*/
849
850
static PyObject *
851
itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
852
                             PyObject *values, PyObject *next)
853
/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
854
0
{
855
0
    teedataobject *tdo;
856
0
    Py_ssize_t i, len;
857
858
0
    itertools_state *state = get_module_state_by_cls(type);
859
0
    assert(type == state->teedataobject_type);
860
861
0
    tdo = (teedataobject *)teedataobject_newinternal(state, it);
862
0
    if (!tdo)
863
0
        return NULL;
864
865
0
    len = PyList_GET_SIZE(values);
866
0
    if (len > LINKCELLS)
867
0
        goto err;
868
0
    for (i=0; i<len; i++) {
869
0
        tdo->values[i] = PyList_GET_ITEM(values, i);
870
0
        Py_INCREF(tdo->values[i]);
871
0
    }
872
    /* len <= LINKCELLS < INT_MAX */
873
0
    tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
874
875
0
    if (len == LINKCELLS) {
876
0
        if (next != Py_None) {
877
0
            if (!Py_IS_TYPE(next, state->teedataobject_type))
878
0
                goto err;
879
0
            assert(tdo->nextlink == NULL);
880
0
            tdo->nextlink = Py_NewRef(next);
881
0
        }
882
0
    } else {
883
0
        if (next != Py_None)
884
0
            goto err; /* shouldn't have a next if we are not full */
885
0
    }
886
0
    return (PyObject*)tdo;
887
888
0
err:
889
0
    Py_XDECREF(tdo);
890
0
    PyErr_SetString(PyExc_ValueError, "Invalid arguments");
891
0
    return NULL;
892
0
}
893
894
static PyType_Slot teedataobject_slots[] = {
895
    {Py_tp_dealloc, teedataobject_dealloc},
896
    {Py_tp_getattro, PyObject_GenericGetAttr},
897
    {Py_tp_doc, (void *)itertools_teedataobject__doc__},
898
    {Py_tp_traverse, teedataobject_traverse},
899
    {Py_tp_clear, teedataobject_clear},
900
    {Py_tp_new, itertools_teedataobject},
901
    {Py_tp_free, PyObject_GC_Del},
902
    {0, NULL},
903
};
904
905
static PyType_Spec teedataobject_spec = {
906
    .name = "itertools._tee_dataobject",
907
    .basicsize = sizeof(teedataobject),
908
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
909
              Py_TPFLAGS_IMMUTABLETYPE),
910
    .slots = teedataobject_slots,
911
};
912
913
914
static PyObject *
915
tee_next(PyObject *op)
916
0
{
917
0
    teeobject *to = teeobject_CAST(op);
918
0
    PyObject *value, *link;
919
920
0
    if (to->index >= LINKCELLS) {
921
0
        link = teedataobject_jumplink(to->state, to->dataobj);
922
0
        if (link == NULL)
923
0
            return NULL;
924
0
        Py_SETREF(to->dataobj, (teedataobject *)link);
925
0
        to->index = 0;
926
0
    }
927
0
    value = teedataobject_getitem(to->dataobj, to->index);
928
0
    if (value == NULL)
929
0
        return NULL;
930
0
    to->index++;
931
0
    return value;
932
0
}
933
934
static int
935
tee_traverse(PyObject *op, visitproc visit, void *arg)
936
0
{
937
0
    teeobject *to = teeobject_CAST(op);
938
0
    Py_VISIT(Py_TYPE(to));
939
0
    Py_VISIT((PyObject *)to->dataobj);
940
0
    return 0;
941
0
}
942
943
static teeobject *
944
tee_copy_impl(teeobject *to)
945
0
{
946
0
    teeobject *newto = PyObject_GC_New(teeobject, Py_TYPE(to));
947
0
    if (newto == NULL) {
948
0
        return NULL;
949
0
    }
950
0
    newto->dataobj = (teedataobject *)Py_NewRef(to->dataobj);
951
0
    newto->index = to->index;
952
0
    newto->weakreflist = NULL;
953
0
    newto->state = to->state;
954
0
    PyObject_GC_Track(newto);
955
0
    return newto;
956
0
}
957
958
static inline PyObject *
959
tee_copy(PyObject *op, PyObject *Py_UNUSED(ignored))
960
0
{
961
0
    teeobject *to = teeobject_CAST(op);
962
0
    return (PyObject *)tee_copy_impl(to);
963
0
}
964
965
PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
966
967
static PyObject *
968
tee_fromiterable(itertools_state *state, PyObject *iterable)
969
0
{
970
0
    teeobject *to;
971
0
    PyObject *it;
972
973
0
    it = PyObject_GetIter(iterable);
974
0
    if (it == NULL)
975
0
        return NULL;
976
0
    if (PyObject_TypeCheck(it, state->tee_type)) {
977
0
        to = tee_copy_impl((teeobject *)it);  // 'it' can be fast casted
978
0
        goto done;
979
0
    }
980
981
0
    PyObject *dataobj = teedataobject_newinternal(state, it);
982
0
    if (!dataobj) {
983
0
        to = NULL;
984
0
        goto done;
985
0
    }
986
0
    to = PyObject_GC_New(teeobject, state->tee_type);
987
0
    if (to == NULL) {
988
0
        Py_DECREF(dataobj);
989
0
        goto done;
990
0
    }
991
0
    to->dataobj = (teedataobject *)dataobj;
992
0
    to->index = 0;
993
0
    to->weakreflist = NULL;
994
0
    to->state = state;
995
0
    PyObject_GC_Track(to);
996
0
done:
997
0
    Py_DECREF(it);
998
0
    return (PyObject *)to;
999
0
}
1000
1001
/*[clinic input]
1002
@classmethod
1003
itertools._tee.__new__
1004
    iterable: object
1005
    /
1006
Iterator wrapped to make it copyable.
1007
[clinic start generated code]*/
1008
1009
static PyObject *
1010
itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
1011
/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
1012
0
{
1013
0
    itertools_state *state = get_module_state_by_cls(type);
1014
0
    return tee_fromiterable(state, iterable);
1015
0
}
1016
1017
static int
1018
tee_clear(PyObject *op)
1019
0
{
1020
0
    teeobject *to = teeobject_CAST(op);
1021
0
    if (to->weakreflist != NULL)
1022
0
        PyObject_ClearWeakRefs(op);
1023
0
    Py_CLEAR(to->dataobj);
1024
0
    return 0;
1025
0
}
1026
1027
static void
1028
tee_dealloc(PyObject *op)
1029
0
{
1030
0
    PyTypeObject *tp = Py_TYPE(op);
1031
0
    PyObject_GC_UnTrack(op);
1032
0
    (void)tee_clear(op);
1033
0
    PyObject_GC_Del(op);
1034
0
    Py_DECREF(tp);
1035
0
}
1036
1037
static PyMethodDef tee_methods[] = {
1038
    {"__copy__", tee_copy, METH_NOARGS, teecopy_doc},
1039
    {NULL,              NULL}           /* sentinel */
1040
};
1041
1042
static PyMemberDef tee_members[] = {
1043
    {"__weaklistoffset__", Py_T_PYSSIZET, offsetof(teeobject, weakreflist), Py_READONLY},
1044
    {NULL},
1045
};
1046
1047
static PyType_Slot tee_slots[] = {
1048
    {Py_tp_dealloc, tee_dealloc},
1049
    {Py_tp_doc, (void *)itertools__tee__doc__},
1050
    {Py_tp_traverse, tee_traverse},
1051
    {Py_tp_clear, tee_clear},
1052
    {Py_tp_iter, PyObject_SelfIter},
1053
    {Py_tp_iternext, tee_next},
1054
    {Py_tp_methods, tee_methods},
1055
    {Py_tp_members, tee_members},
1056
    {Py_tp_new, itertools__tee},
1057
    {Py_tp_free, PyObject_GC_Del},
1058
    {0, NULL},
1059
};
1060
1061
static PyType_Spec tee_spec = {
1062
    .name = "itertools._tee",
1063
    .basicsize = sizeof(teeobject),
1064
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1065
              Py_TPFLAGS_IMMUTABLETYPE),
1066
    .slots = tee_slots,
1067
};
1068
1069
/*[clinic input]
1070
itertools.tee
1071
    iterable: object
1072
    n: Py_ssize_t(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
0
#define chainobject_CAST(op)    ((chainobject *)(op))
1804
1805
static PyObject *
1806
chain_new_internal(PyTypeObject *type, PyObject *source)
1807
0
{
1808
0
    chainobject *lz;
1809
1810
0
    lz = (chainobject *)type->tp_alloc(type, 0);
1811
0
    if (lz == NULL) {
1812
0
        Py_DECREF(source);
1813
0
        return NULL;
1814
0
    }
1815
1816
0
    lz->source = source;
1817
0
    lz->active = NULL;
1818
0
    return (PyObject *)lz;
1819
0
}
1820
1821
static PyObject *
1822
chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1823
0
{
1824
0
    PyObject *source;
1825
1826
0
    itertools_state *state = find_state_by_type(type);
1827
0
    PyTypeObject *chain_type = state->chain_type;
1828
0
    if ((type == chain_type || type->tp_init == chain_type->tp_init) &&
1829
0
        !_PyArg_NoKeywords("chain", kwds))
1830
0
        return NULL;
1831
1832
0
    source = PyObject_GetIter(args);
1833
0
    if (source == NULL)
1834
0
        return NULL;
1835
1836
0
    return chain_new_internal(type, source);
1837
0
}
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
0
{
1852
0
    PyObject *source;
1853
1854
0
    source = PyObject_GetIter(arg);
1855
0
    if (source == NULL)
1856
0
        return NULL;
1857
1858
0
    return chain_new_internal(type, source);
1859
0
}
1860
1861
static void
1862
chain_dealloc(PyObject *op)
1863
0
{
1864
0
    chainobject *lz = chainobject_CAST(op);
1865
0
    PyTypeObject *tp = Py_TYPE(lz);
1866
0
    PyObject_GC_UnTrack(lz);
1867
0
    Py_XDECREF(lz->active);
1868
0
    Py_XDECREF(lz->source);
1869
0
    tp->tp_free(lz);
1870
0
    Py_DECREF(tp);
1871
0
}
1872
1873
static int
1874
chain_traverse(PyObject *op, visitproc visit, void *arg)
1875
0
{
1876
0
    chainobject *lz = chainobject_CAST(op);
1877
0
    Py_VISIT(Py_TYPE(lz));
1878
0
    Py_VISIT(lz->source);
1879
0
    Py_VISIT(lz->active);
1880
0
    return 0;
1881
0
}
1882
1883
static inline PyObject *
1884
chain_next_lock_held(PyObject *op)
1885
0
{
1886
0
    chainobject *lz = chainobject_CAST(op);
1887
0
    PyObject *item;
1888
1889
    /* lz->source is the iterator of iterables. If it's NULL, we've already
1890
     * consumed them all. lz->active is the current iterator. If it's NULL,
1891
     * we should grab a new one from lz->source. */
1892
0
    while (lz->source != NULL) {
1893
0
        if (lz->active == NULL) {
1894
0
            PyObject *iterable = PyIter_Next(lz->source);
1895
0
            if (iterable == NULL) {
1896
0
                Py_CLEAR(lz->source);
1897
0
                return NULL;            /* no more input sources */
1898
0
            }
1899
0
            lz->active = PyObject_GetIter(iterable);
1900
0
            Py_DECREF(iterable);
1901
0
            if (lz->active == NULL) {
1902
0
                Py_CLEAR(lz->source);
1903
0
                return NULL;            /* input not iterable */
1904
0
            }
1905
0
        }
1906
0
        item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1907
0
        if (item != NULL)
1908
0
            return item;
1909
0
        if (PyErr_Occurred()) {
1910
0
            if (PyErr_ExceptionMatches(PyExc_StopIteration))
1911
0
                PyErr_Clear();
1912
0
            else
1913
0
                return NULL;            /* input raised an exception */
1914
0
        }
1915
        /* lz->active is consumed, try with the next iterable. */
1916
0
        Py_CLEAR(lz->active);
1917
0
    }
1918
    /* Everything had been consumed already. */
1919
0
    return NULL;
1920
0
}
1921
1922
static PyObject *
1923
chain_next(PyObject *op)
1924
0
{
1925
0
    PyObject *result;
1926
0
    Py_BEGIN_CRITICAL_SECTION(op);
1927
0
    result = chain_next_lock_held(op);
1928
0
    Py_END_CRITICAL_SECTION()
1929
0
    return result;
1930
0
}
1931
1932
PyDoc_STRVAR(chain_doc,
1933
"chain(*iterables)\n\
1934
--\n\
1935
\n\
1936
Return a chain object whose .__next__() method returns elements from the\n\
1937
first iterable until it is exhausted, then elements from the next\n\
1938
iterable, until all of the iterables are exhausted.");
1939
1940
static PyMethodDef chain_methods[] = {
1941
    ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
1942
    {"__class_getitem__",    Py_GenericAlias,
1943
    METH_O|METH_CLASS,       PyDoc_STR("See PEP 585")},
1944
    {NULL,              NULL}           /* sentinel */
1945
};
1946
1947
static PyType_Slot chain_slots[] = {
1948
    {Py_tp_dealloc, chain_dealloc},
1949
    {Py_tp_getattro, PyObject_GenericGetAttr},
1950
    {Py_tp_doc, (void *)chain_doc},
1951
    {Py_tp_traverse, chain_traverse},
1952
    {Py_tp_iter, PyObject_SelfIter},
1953
    {Py_tp_iternext, chain_next},
1954
    {Py_tp_methods, chain_methods},
1955
    {Py_tp_new, chain_new},
1956
    {Py_tp_free, PyObject_GC_Del},
1957
    {0, NULL},
1958
};
1959
1960
static PyType_Spec chain_spec = {
1961
    .name = "itertools.chain",
1962
    .basicsize = sizeof(chainobject),
1963
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1964
              Py_TPFLAGS_IMMUTABLETYPE),
1965
    .slots = chain_slots,
1966
};
1967
1968
1969
/* product object ************************************************************/
1970
1971
typedef struct {
1972
    PyObject_HEAD
1973
    PyObject *pools;        /* tuple of pool tuples */
1974
    Py_ssize_t *indices;    /* one index per pool */
1975
    PyObject *result;       /* most recently returned result tuple */
1976
    int stopped;            /* set to 1 when the iterator is exhausted */
1977
} productobject;
1978
1979
336
#define productobject_CAST(op)  ((productobject *)(op))
1980
1981
static PyObject *
1982
product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1983
66
{
1984
66
    productobject *lz;
1985
66
    Py_ssize_t nargs, npools, repeat=1;
1986
66
    PyObject *pools = NULL;
1987
66
    Py_ssize_t *indices = NULL;
1988
66
    Py_ssize_t i;
1989
1990
66
    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
66
    assert(PyTuple_CheckExact(args));
2009
66
    if (repeat == 0) {
2010
0
        nargs = 0;
2011
66
    } else {
2012
66
        nargs = PyTuple_GET_SIZE(args);
2013
66
        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
66
    }
2018
66
    npools = nargs * repeat;
2019
2020
66
    indices = PyMem_New(Py_ssize_t, npools);
2021
66
    if (indices == NULL) {
2022
0
        PyErr_NoMemory();
2023
0
        goto error;
2024
0
    }
2025
2026
66
    pools = PyTuple_New(npools);
2027
66
    if (pools == NULL)
2028
0
        goto error;
2029
2030
168
    for (i=0; i < nargs ; ++i) {
2031
102
        PyObject *item = PyTuple_GET_ITEM(args, i);
2032
102
        PyObject *pool = PySequence_Tuple(item);
2033
102
        if (pool == NULL)
2034
0
            goto error;
2035
102
        PyTuple_SET_ITEM(pools, i, pool);
2036
102
        indices[i] = 0;
2037
102
    }
2038
66
    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
66
    lz = (productobject *)type->tp_alloc(type, 0);
2047
66
    if (lz == NULL)
2048
0
        goto error;
2049
2050
66
    lz->pools = pools;
2051
66
    lz->indices = indices;
2052
66
    lz->result = NULL;
2053
66
    lz->stopped = 0;
2054
2055
66
    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
66
}
2063
2064
static void
2065
product_dealloc(PyObject *op)
2066
66
{
2067
66
    productobject *lz = productobject_CAST(op);
2068
66
    PyTypeObject *tp = Py_TYPE(lz);
2069
66
    PyObject_GC_UnTrack(lz);
2070
66
    Py_XDECREF(lz->pools);
2071
66
    Py_XDECREF(lz->result);
2072
66
    PyMem_Free(lz->indices);
2073
66
    tp->tp_free(lz);
2074
66
    Py_DECREF(tp);
2075
66
}
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
270
{
2101
270
    productobject *lz = productobject_CAST(op);
2102
270
    PyObject *pool;
2103
270
    PyObject *elem;
2104
270
    PyObject *oldelem;
2105
270
    PyObject *pools = lz->pools;
2106
270
    PyObject *result = lz->result;
2107
270
    Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2108
270
    Py_ssize_t i;
2109
2110
270
    if (lz->stopped)
2111
0
        return NULL;
2112
2113
270
    if (result == NULL) {
2114
        /* On the first pass, return an initial tuple filled with the
2115
           first element from each pool. */
2116
66
        result = PyTuple_New(npools);
2117
66
        if (result == NULL)
2118
0
            goto empty;
2119
66
        lz->result = result;
2120
168
        for (i=0; i < npools; i++) {
2121
102
            pool = PyTuple_GET_ITEM(pools, i);
2122
102
            if (PyTuple_GET_SIZE(pool) == 0)
2123
0
                goto empty;
2124
102
            elem = PyTuple_GET_ITEM(pool, 0);
2125
102
            Py_INCREF(elem);
2126
102
            PyTuple_SET_ITEM(result, i, elem);
2127
102
        }
2128
204
    } else {
2129
204
        Py_ssize_t *indices = lz->indices;
2130
2131
        /* Copy the previous result tuple or re-use it if available */
2132
204
        if (Py_REFCNT(result) > 1) {
2133
204
            PyObject *old_result = result;
2134
204
            result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
2135
204
            if (result == NULL)
2136
0
                goto empty;
2137
204
            lz->result = result;
2138
204
            Py_DECREF(old_result);
2139
204
        }
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
204
        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
342
        for (i=npools-1 ; i >= 0 ; i--) {
2151
276
            pool = PyTuple_GET_ITEM(pools, i);
2152
276
            indices[i]++;
2153
276
            if (indices[i] == PyTuple_GET_SIZE(pool)) {
2154
                /* Roll-over and advance to next pool */
2155
138
                indices[i] = 0;
2156
138
                elem = PyTuple_GET_ITEM(pool, 0);
2157
138
                Py_INCREF(elem);
2158
138
                oldelem = PyTuple_GET_ITEM(result, i);
2159
138
                PyTuple_SET_ITEM(result, i, elem);
2160
138
                Py_DECREF(oldelem);
2161
138
            } else {
2162
                /* No rollover. Just increment and stop here. */
2163
138
                elem = PyTuple_GET_ITEM(pool, indices[i]);
2164
138
                Py_INCREF(elem);
2165
138
                oldelem = PyTuple_GET_ITEM(result, i);
2166
138
                PyTuple_SET_ITEM(result, i, elem);
2167
138
                Py_DECREF(oldelem);
2168
138
                break;
2169
138
            }
2170
276
        }
2171
2172
        /* If i is negative, then the indices have all rolled-over
2173
           and we're done. */
2174
204
        if (i < 0)
2175
66
            goto empty;
2176
204
    }
2177
2178
204
    return Py_NewRef(result);
2179
2180
66
empty:
2181
66
    lz->stopped = 1;
2182
66
    return NULL;
2183
270
}
2184
2185
static PyObject *
2186
product_next(PyObject *op)
2187
270
{
2188
270
    PyObject *result;
2189
270
    Py_BEGIN_CRITICAL_SECTION(op);
2190
270
    result = product_next_lock_held(op);
2191
270
    Py_END_CRITICAL_SECTION()
2192
270
    return result;
2193
270
}
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 (Py_REFCNT(result) > 1) {
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(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 (Py_REFCNT(result) > 1) {
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 PyMethodDef cwr_methods[] = {
2670
    {"__sizeof__", cwr_sizeof, METH_NOARGS, sizeof_doc},
2671
    {NULL,              NULL}   /* sentinel */
2672
};
2673
2674
static PyType_Slot cwr_slots[] = {
2675
    {Py_tp_dealloc, cwr_dealloc},
2676
    {Py_tp_getattro, PyObject_GenericGetAttr},
2677
    {Py_tp_doc, (void *)itertools_combinations_with_replacement__doc__},
2678
    {Py_tp_traverse, cwr_traverse},
2679
    {Py_tp_iter, PyObject_SelfIter},
2680
    {Py_tp_iternext, cwr_next},
2681
    {Py_tp_methods, cwr_methods},
2682
    {Py_tp_new, itertools_combinations_with_replacement},
2683
    {Py_tp_free, PyObject_GC_Del},
2684
    {0, NULL},
2685
};
2686
2687
static PyType_Spec cwr_spec = {
2688
    .name = "itertools.combinations_with_replacement",
2689
    .basicsize = sizeof(cwrobject),
2690
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2691
              Py_TPFLAGS_IMMUTABLETYPE),
2692
    .slots = cwr_slots,
2693
};
2694
2695
2696
/* permutations object ********************************************************
2697
2698
def permutations(iterable, r=None):
2699
    # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
2700
    # permutations(range(3)) --> 012 021 102 120 201 210
2701
    pool = tuple(iterable)
2702
    n = len(pool)
2703
    r = n if r is None else r
2704
    if r > n:
2705
        return
2706
    indices = list(range(n))
2707
    cycles = list(range(n, n-r, -1))
2708
    yield tuple(pool[i] for i in indices[:r])
2709
    while n:
2710
        for i in reversed(range(r)):
2711
            cycles[i] -= 1
2712
            if cycles[i] == 0:
2713
                indices[i:] = indices[i+1:] + indices[i:i+1]
2714
                cycles[i] = n - i
2715
            else:
2716
                j = cycles[i]
2717
                indices[i], indices[-j] = indices[-j], indices[i]
2718
                yield tuple(pool[i] for i in indices[:r])
2719
                break
2720
        else:
2721
            return
2722
*/
2723
2724
typedef struct {
2725
    PyObject_HEAD
2726
    PyObject *pool;         /* input converted to a tuple */
2727
    Py_ssize_t *indices;    /* one index per element in the pool */
2728
    Py_ssize_t *cycles;     /* one rollover counter per element in the result */
2729
    PyObject *result;       /* most recently returned result tuple */
2730
    Py_ssize_t r;           /* size of result tuple */
2731
    int stopped;            /* set to 1 when the iterator is exhausted */
2732
} permutationsobject;
2733
2734
162
#define permutationsobject_CAST(op) ((permutationsobject *)(op))
2735
2736
/*[clinic input]
2737
@classmethod
2738
itertools.permutations.__new__
2739
    iterable: object
2740
    r as robj: object = None
2741
Return successive r-length permutations of elements in the iterable.
2742
2743
permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
2744
[clinic start generated code]*/
2745
2746
static PyObject *
2747
itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
2748
                            PyObject *robj)
2749
/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
2750
48
{
2751
48
    permutationsobject *po;
2752
48
    Py_ssize_t n;
2753
48
    Py_ssize_t r;
2754
48
    PyObject *pool = NULL;
2755
48
    Py_ssize_t *indices = NULL;
2756
48
    Py_ssize_t *cycles = NULL;
2757
48
    Py_ssize_t i;
2758
2759
48
    pool = PySequence_Tuple(iterable);
2760
48
    if (pool == NULL)
2761
0
        goto error;
2762
48
    n = PyTuple_GET_SIZE(pool);
2763
2764
48
    r = n;
2765
48
    if (robj != Py_None) {
2766
0
        if (!PyLong_Check(robj)) {
2767
0
            PyErr_SetString(PyExc_TypeError, "Expected int as r");
2768
0
            goto error;
2769
0
        }
2770
0
        r = PyLong_AsSsize_t(robj);
2771
0
        if (r == -1 && PyErr_Occurred())
2772
0
            goto error;
2773
0
    }
2774
48
    if (r < 0) {
2775
0
        PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2776
0
        goto error;
2777
0
    }
2778
2779
48
    indices = PyMem_New(Py_ssize_t, n);
2780
48
    cycles = PyMem_New(Py_ssize_t, r);
2781
48
    if (indices == NULL || cycles == NULL) {
2782
0
        PyErr_NoMemory();
2783
0
        goto error;
2784
0
    }
2785
2786
114
    for (i=0 ; i<n ; i++)
2787
66
        indices[i] = i;
2788
114
    for (i=0 ; i<r ; i++)
2789
66
        cycles[i] = n - i;
2790
2791
    /* create permutationsobject structure */
2792
48
    po = (permutationsobject *)type->tp_alloc(type, 0);
2793
48
    if (po == NULL)
2794
0
        goto error;
2795
2796
48
    po->pool = pool;
2797
48
    po->indices = indices;
2798
48
    po->cycles = cycles;
2799
48
    po->result = NULL;
2800
48
    po->r = r;
2801
48
    po->stopped = r > n ? 1 : 0;
2802
2803
48
    return (PyObject *)po;
2804
2805
0
error:
2806
0
    if (indices != NULL)
2807
0
        PyMem_Free(indices);
2808
0
    if (cycles != NULL)
2809
0
        PyMem_Free(cycles);
2810
0
    Py_XDECREF(pool);
2811
0
    return NULL;
2812
48
}
2813
2814
static void
2815
permutations_dealloc(PyObject *op)
2816
48
{
2817
48
    permutationsobject *po = permutationsobject_CAST(op);
2818
48
    PyTypeObject *tp = Py_TYPE(po);
2819
48
    PyObject_GC_UnTrack(po);
2820
48
    Py_XDECREF(po->pool);
2821
48
    Py_XDECREF(po->result);
2822
48
    PyMem_Free(po->indices);
2823
48
    PyMem_Free(po->cycles);
2824
48
    tp->tp_free(po);
2825
48
    Py_DECREF(tp);
2826
48
}
2827
2828
static PyObject *
2829
permutations_sizeof(PyObject *op, PyObject *Py_UNUSED(args))
2830
0
{
2831
0
    permutationsobject *po = permutationsobject_CAST(op);
2832
0
    size_t res = _PyObject_SIZE(Py_TYPE(po));
2833
0
    res += (size_t)PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
2834
0
    res += (size_t)po->r * sizeof(Py_ssize_t);
2835
0
    return PyLong_FromSize_t(res);
2836
0
}
2837
2838
static int
2839
permutations_traverse(PyObject *op, visitproc visit, void *arg)
2840
0
{
2841
0
    permutationsobject *po = permutationsobject_CAST(op);
2842
0
    Py_VISIT(Py_TYPE(po));
2843
0
    Py_VISIT(po->pool);
2844
0
    Py_VISIT(po->result);
2845
0
    return 0;
2846
0
}
2847
2848
static PyObject *
2849
permutations_next(PyObject *op)
2850
114
{
2851
114
    permutationsobject *po = permutationsobject_CAST(op);
2852
114
    PyObject *elem;
2853
114
    PyObject *oldelem;
2854
114
    PyObject *pool = po->pool;
2855
114
    Py_ssize_t *indices = po->indices;
2856
114
    Py_ssize_t *cycles = po->cycles;
2857
114
    PyObject *result = po->result;
2858
114
    Py_ssize_t n = PyTuple_GET_SIZE(pool);
2859
114
    Py_ssize_t r = po->r;
2860
114
    Py_ssize_t i, j, k, index;
2861
2862
114
    if (po->stopped)
2863
0
        return NULL;
2864
2865
114
    if (result == NULL) {
2866
        /* On the first pass, initialize result tuple using the indices */
2867
48
        result = PyTuple_New(r);
2868
48
        if (result == NULL)
2869
0
            goto empty;
2870
48
        po->result = result;
2871
114
        for (i=0; i<r ; i++) {
2872
66
            index = indices[i];
2873
66
            elem = PyTuple_GET_ITEM(pool, index);
2874
66
            Py_INCREF(elem);
2875
66
            PyTuple_SET_ITEM(result, i, elem);
2876
66
        }
2877
66
    } else {
2878
66
        if (n == 0)
2879
0
            goto empty;
2880
2881
        /* Copy the previous result tuple or re-use it if available */
2882
66
        if (Py_REFCNT(result) > 1) {
2883
66
            PyObject *old_result = result;
2884
66
            result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2885
66
            if (result == NULL)
2886
0
                goto empty;
2887
66
            po->result = result;
2888
66
            Py_DECREF(old_result);
2889
66
        }
2890
        // bpo-42536: The GC may have untracked this result tuple. Since we're
2891
        // recycling it, make sure it's tracked again:
2892
0
        else {
2893
0
            _PyTuple_Recycle(result);
2894
0
        }
2895
        /* Now, we've got the only copy so we can update it in-place */
2896
66
        assert(r == 0 || Py_REFCNT(result) == 1);
2897
2898
        /* Decrement rightmost cycle, moving leftward upon zero rollover */
2899
150
        for (i=r-1 ; i>=0 ; i--) {
2900
102
            cycles[i] -= 1;
2901
102
            if (cycles[i] == 0) {
2902
                /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2903
84
                index = indices[i];
2904
102
                for (j=i ; j<n-1 ; j++)
2905
18
                    indices[j] = indices[j+1];
2906
84
                indices[n-1] = index;
2907
84
                cycles[i] = n - i;
2908
84
            } else {
2909
18
                j = cycles[i];
2910
18
                index = indices[i];
2911
18
                indices[i] = indices[n-j];
2912
18
                indices[n-j] = index;
2913
2914
54
                for (k=i; k<r ; k++) {
2915
                    /* start with i, the leftmost element that changed */
2916
                    /* yield tuple(pool[k] for k in indices[:r]) */
2917
36
                    index = indices[k];
2918
36
                    elem = PyTuple_GET_ITEM(pool, index);
2919
36
                    Py_INCREF(elem);
2920
36
                    oldelem = PyTuple_GET_ITEM(result, k);
2921
36
                    PyTuple_SET_ITEM(result, k, elem);
2922
36
                    Py_DECREF(oldelem);
2923
36
                }
2924
18
                break;
2925
18
            }
2926
102
        }
2927
        /* If i is negative, then the cycles have all
2928
           rolled-over and we're done. */
2929
66
        if (i < 0)
2930
48
            goto empty;
2931
66
    }
2932
66
    return Py_NewRef(result);
2933
2934
48
empty:
2935
48
    po->stopped = 1;
2936
48
    return NULL;
2937
114
}
2938
2939
static PyMethodDef permuations_methods[] = {
2940
    {"__sizeof__", permutations_sizeof, METH_NOARGS, sizeof_doc},
2941
    {NULL,              NULL}   /* sentinel */
2942
};
2943
2944
static PyType_Slot permutations_slots[] = {
2945
    {Py_tp_dealloc, permutations_dealloc},
2946
    {Py_tp_getattro, PyObject_GenericGetAttr},
2947
    {Py_tp_doc, (void *)itertools_permutations__doc__},
2948
    {Py_tp_traverse, permutations_traverse},
2949
    {Py_tp_iter, PyObject_SelfIter},
2950
    {Py_tp_iternext, permutations_next},
2951
    {Py_tp_methods, permuations_methods},
2952
    {Py_tp_new, itertools_permutations},
2953
    {Py_tp_free, PyObject_GC_Del},
2954
    {0, NULL},
2955
};
2956
2957
static PyType_Spec permutations_spec = {
2958
    .name = "itertools.permutations",
2959
    .basicsize = sizeof(permutationsobject),
2960
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2961
              Py_TPFLAGS_IMMUTABLETYPE),
2962
    .slots = permutations_slots,
2963
};
2964
2965
2966
/* accumulate object ********************************************************/
2967
2968
typedef struct {
2969
    PyObject_HEAD
2970
    PyObject *total;
2971
    PyObject *it;
2972
    PyObject *binop;
2973
    PyObject *initial;
2974
    itertools_state *state;
2975
} accumulateobject;
2976
2977
0
#define accumulateobject_CAST(op)   ((accumulateobject *)(op))
2978
2979
/*[clinic input]
2980
@classmethod
2981
itertools.accumulate.__new__
2982
    iterable: object
2983
    func as binop: object = None
2984
    *
2985
    initial: object = None
2986
Return series of accumulated sums (or other binary function results).
2987
[clinic start generated code]*/
2988
2989
static PyObject *
2990
itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
2991
                          PyObject *binop, PyObject *initial)
2992
/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
2993
0
{
2994
0
    PyObject *it;
2995
0
    accumulateobject *lz;
2996
2997
    /* Get iterator. */
2998
0
    it = PyObject_GetIter(iterable);
2999
0
    if (it == NULL)
3000
0
        return NULL;
3001
3002
    /* create accumulateobject structure */
3003
0
    lz = (accumulateobject *)type->tp_alloc(type, 0);
3004
0
    if (lz == NULL) {
3005
0
        Py_DECREF(it);
3006
0
        return NULL;
3007
0
    }
3008
3009
0
    if (binop != Py_None) {
3010
0
        lz->binop = Py_XNewRef(binop);
3011
0
    }
3012
0
    lz->total = NULL;
3013
0
    lz->it = it;
3014
0
    lz->initial = Py_XNewRef(initial);
3015
0
    lz->state = find_state_by_type(type);
3016
0
    return (PyObject *)lz;
3017
0
}
3018
3019
static void
3020
accumulate_dealloc(PyObject *op)
3021
0
{
3022
0
    accumulateobject *lz = accumulateobject_CAST(op);
3023
0
    PyTypeObject *tp = Py_TYPE(lz);
3024
0
    PyObject_GC_UnTrack(lz);
3025
0
    Py_XDECREF(lz->binop);
3026
0
    Py_XDECREF(lz->total);
3027
0
    Py_XDECREF(lz->it);
3028
0
    Py_XDECREF(lz->initial);
3029
0
    tp->tp_free(lz);
3030
0
    Py_DECREF(tp);
3031
0
}
3032
3033
static int
3034
accumulate_traverse(PyObject *op, visitproc visit, void *arg)
3035
0
{
3036
0
    accumulateobject *lz = accumulateobject_CAST(op);
3037
0
    Py_VISIT(Py_TYPE(lz));
3038
0
    Py_VISIT(lz->binop);
3039
0
    Py_VISIT(lz->it);
3040
0
    Py_VISIT(lz->total);
3041
0
    Py_VISIT(lz->initial);
3042
0
    return 0;
3043
0
}
3044
3045
static PyObject *
3046
accumulate_next(PyObject *op)
3047
0
{
3048
0
    accumulateobject *lz = accumulateobject_CAST(op);
3049
0
    PyObject *val, *newtotal;
3050
3051
0
    if (lz->initial != Py_None) {
3052
0
        lz->total = lz->initial;
3053
0
        lz->initial = Py_NewRef(Py_None);
3054
0
        return Py_NewRef(lz->total);
3055
0
    }
3056
0
    val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
3057
0
    if (val == NULL)
3058
0
        return NULL;
3059
3060
0
    if (lz->total == NULL) {
3061
0
        lz->total = Py_NewRef(val);
3062
0
        return lz->total;
3063
0
    }
3064
3065
0
    if (lz->binop == NULL)
3066
0
        newtotal = PyNumber_Add(lz->total, val);
3067
0
    else
3068
0
        newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
3069
0
    Py_DECREF(val);
3070
0
    if (newtotal == NULL)
3071
0
        return NULL;
3072
3073
0
    Py_INCREF(newtotal);
3074
0
    Py_SETREF(lz->total, newtotal);
3075
0
    return newtotal;
3076
0
}
3077
3078
static PyType_Slot accumulate_slots[] = {
3079
    {Py_tp_dealloc, accumulate_dealloc},
3080
    {Py_tp_getattro, PyObject_GenericGetAttr},
3081
    {Py_tp_doc, (void *)itertools_accumulate__doc__},
3082
    {Py_tp_traverse, accumulate_traverse},
3083
    {Py_tp_iter, PyObject_SelfIter},
3084
    {Py_tp_iternext, accumulate_next},
3085
    {Py_tp_new, itertools_accumulate},
3086
    {Py_tp_free, PyObject_GC_Del},
3087
    {0, NULL},
3088
};
3089
3090
static PyType_Spec accumulate_spec = {
3091
    .name = "itertools.accumulate",
3092
    .basicsize = sizeof(accumulateobject),
3093
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3094
              Py_TPFLAGS_IMMUTABLETYPE),
3095
    .slots = accumulate_slots,
3096
};
3097
3098
3099
/* compress object ************************************************************/
3100
3101
/* Equivalent to:
3102
3103
    def compress(data, selectors):
3104
        "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3105
        return (d for d, s in zip(data, selectors) if s)
3106
*/
3107
3108
typedef struct {
3109
    PyObject_HEAD
3110
    PyObject *data;
3111
    PyObject *selectors;
3112
} compressobject;
3113
3114
0
#define compressobject_CAST(op) ((compressobject *)(op))
3115
3116
/*[clinic input]
3117
@classmethod
3118
itertools.compress.__new__
3119
    data as seq1: object
3120
    selectors as seq2: object
3121
Return data elements corresponding to true selector elements.
3122
3123
Forms a shorter iterator from selected data elements using the selectors to
3124
choose the data elements.
3125
[clinic start generated code]*/
3126
3127
static PyObject *
3128
itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3129
/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
3130
0
{
3131
0
    PyObject *data=NULL, *selectors=NULL;
3132
0
    compressobject *lz;
3133
3134
0
    data = PyObject_GetIter(seq1);
3135
0
    if (data == NULL)
3136
0
        goto fail;
3137
0
    selectors = PyObject_GetIter(seq2);
3138
0
    if (selectors == NULL)
3139
0
        goto fail;
3140
3141
    /* create compressobject structure */
3142
0
    lz = (compressobject *)type->tp_alloc(type, 0);
3143
0
    if (lz == NULL)
3144
0
        goto fail;
3145
0
    lz->data = data;
3146
0
    lz->selectors = selectors;
3147
0
    return (PyObject *)lz;
3148
3149
0
fail:
3150
0
    Py_XDECREF(data);
3151
0
    Py_XDECREF(selectors);
3152
0
    return NULL;
3153
0
}
3154
3155
static void
3156
compress_dealloc(PyObject *op)
3157
0
{
3158
0
    compressobject *lz = compressobject_CAST(op);
3159
0
    PyTypeObject *tp = Py_TYPE(lz);
3160
0
    PyObject_GC_UnTrack(lz);
3161
0
    Py_XDECREF(lz->data);
3162
0
    Py_XDECREF(lz->selectors);
3163
0
    tp->tp_free(lz);
3164
0
    Py_DECREF(tp);
3165
0
}
3166
3167
static int
3168
compress_traverse(PyObject *op, visitproc visit, void *arg)
3169
0
{
3170
0
    compressobject *lz = compressobject_CAST(op);
3171
0
    Py_VISIT(Py_TYPE(lz));
3172
0
    Py_VISIT(lz->data);
3173
0
    Py_VISIT(lz->selectors);
3174
0
    return 0;
3175
0
}
3176
3177
static PyObject *
3178
compress_next(PyObject *op)
3179
0
{
3180
0
    compressobject *lz = compressobject_CAST(op);
3181
0
    PyObject *data = lz->data, *selectors = lz->selectors;
3182
0
    PyObject *datum, *selector;
3183
0
    PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3184
0
    PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3185
0
    int ok;
3186
3187
0
    while (1) {
3188
        /* Steps:  get datum, get selector, evaluate selector.
3189
           Order is important (to match the pure python version
3190
           in terms of which input gets a chance to raise an
3191
           exception first).
3192
        */
3193
3194
0
        datum = datanext(data);
3195
0
        if (datum == NULL)
3196
0
            return NULL;
3197
3198
0
        selector = selectornext(selectors);
3199
0
        if (selector == NULL) {
3200
0
            Py_DECREF(datum);
3201
0
            return NULL;
3202
0
        }
3203
3204
0
        ok = PyObject_IsTrue(selector);
3205
0
        Py_DECREF(selector);
3206
0
        if (ok > 0)
3207
0
            return datum;
3208
0
        Py_DECREF(datum);
3209
0
        if (ok < 0)
3210
0
            return NULL;
3211
0
    }
3212
0
}
3213
3214
static PyType_Slot compress_slots[] = {
3215
    {Py_tp_dealloc, compress_dealloc},
3216
    {Py_tp_getattro, PyObject_GenericGetAttr},
3217
    {Py_tp_doc, (void *)itertools_compress__doc__},
3218
    {Py_tp_traverse, compress_traverse},
3219
    {Py_tp_iter, PyObject_SelfIter},
3220
    {Py_tp_iternext, compress_next},
3221
    {Py_tp_new, itertools_compress},
3222
    {Py_tp_free, PyObject_GC_Del},
3223
    {0, NULL},
3224
};
3225
3226
static PyType_Spec compress_spec = {
3227
    .name = "itertools.compress",
3228
    .basicsize = sizeof(compressobject),
3229
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3230
              Py_TPFLAGS_IMMUTABLETYPE),
3231
    .slots = compress_slots,
3232
};
3233
3234
3235
/* filterfalse object ************************************************************/
3236
3237
typedef struct {
3238
    PyObject_HEAD
3239
    PyObject *func;
3240
    PyObject *it;
3241
} filterfalseobject;
3242
3243
0
#define filterfalseobject_CAST(op)  ((filterfalseobject *)(op))
3244
3245
/*[clinic input]
3246
@classmethod
3247
itertools.filterfalse.__new__
3248
    function as func: object
3249
    iterable as seq: object
3250
    /
3251
Return those items of iterable for which function(item) is false.
3252
3253
If function is None, return the items that are false.
3254
[clinic start generated code]*/
3255
3256
static PyObject *
3257
itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3258
/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
3259
0
{
3260
0
    PyObject *it;
3261
0
    filterfalseobject *lz;
3262
3263
    /* Get iterator. */
3264
0
    it = PyObject_GetIter(seq);
3265
0
    if (it == NULL)
3266
0
        return NULL;
3267
3268
    /* create filterfalseobject structure */
3269
0
    lz = (filterfalseobject *)type->tp_alloc(type, 0);
3270
0
    if (lz == NULL) {
3271
0
        Py_DECREF(it);
3272
0
        return NULL;
3273
0
    }
3274
0
    lz->func = Py_NewRef(func);
3275
0
    lz->it = it;
3276
3277
0
    return (PyObject *)lz;
3278
0
}
3279
3280
static void
3281
filterfalse_dealloc(PyObject *op)
3282
0
{
3283
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3284
0
    PyTypeObject *tp = Py_TYPE(lz);
3285
0
    PyObject_GC_UnTrack(lz);
3286
0
    Py_XDECREF(lz->func);
3287
0
    Py_XDECREF(lz->it);
3288
0
    tp->tp_free(lz);
3289
0
    Py_DECREF(tp);
3290
0
}
3291
3292
static int
3293
filterfalse_traverse(PyObject *op, visitproc visit, void *arg)
3294
0
{
3295
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3296
0
    Py_VISIT(Py_TYPE(lz));
3297
0
    Py_VISIT(lz->it);
3298
0
    Py_VISIT(lz->func);
3299
0
    return 0;
3300
0
}
3301
3302
static PyObject *
3303
filterfalse_next(PyObject *op)
3304
0
{
3305
0
    filterfalseobject *lz = filterfalseobject_CAST(op);
3306
0
    PyObject *item;
3307
0
    PyObject *it = lz->it;
3308
0
    long ok;
3309
0
    PyObject *(*iternext)(PyObject *);
3310
3311
0
    iternext = *Py_TYPE(it)->tp_iternext;
3312
0
    for (;;) {
3313
0
        item = iternext(it);
3314
0
        if (item == NULL)
3315
0
            return NULL;
3316
3317
0
        if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3318
0
            ok = PyObject_IsTrue(item);
3319
0
        } else {
3320
0
            PyObject *good;
3321
0
            good = PyObject_CallOneArg(lz->func, item);
3322
0
            if (good == NULL) {
3323
0
                Py_DECREF(item);
3324
0
                return NULL;
3325
0
            }
3326
0
            ok = PyObject_IsTrue(good);
3327
0
            Py_DECREF(good);
3328
0
        }
3329
0
        if (ok == 0)
3330
0
            return item;
3331
0
        Py_DECREF(item);
3332
0
        if (ok < 0)
3333
0
            return NULL;
3334
0
    }
3335
0
}
3336
3337
static PyType_Slot filterfalse_slots[] = {
3338
    {Py_tp_dealloc, filterfalse_dealloc},
3339
    {Py_tp_getattro, PyObject_GenericGetAttr},
3340
    {Py_tp_doc, (void *)itertools_filterfalse__doc__},
3341
    {Py_tp_traverse, filterfalse_traverse},
3342
    {Py_tp_iter, PyObject_SelfIter},
3343
    {Py_tp_iternext, filterfalse_next},
3344
    {Py_tp_new, itertools_filterfalse},
3345
    {Py_tp_free, PyObject_GC_Del},
3346
    {0, NULL},
3347
};
3348
3349
static PyType_Spec filterfalse_spec = {
3350
    .name = "itertools.filterfalse",
3351
    .basicsize = sizeof(filterfalseobject),
3352
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3353
              Py_TPFLAGS_IMMUTABLETYPE),
3354
    .slots = filterfalse_slots,
3355
};
3356
3357
3358
/* count object ************************************************************/
3359
3360
typedef struct {
3361
    PyObject_HEAD
3362
    Py_ssize_t cnt;
3363
    PyObject *long_cnt;
3364
    PyObject *long_step;
3365
} countobject;
3366
3367
4.26k
#define countobject_CAST(op)    ((countobject *)(op))
3368
3369
/* Counting logic and invariants:
3370
3371
fast_mode:  when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3372
3373
    assert(long_cnt == NULL && long_step==PyLong(1));
3374
    Advances with:  cnt += 1
3375
    When count hits PY_SSIZE_T_MAX, switch to slow_mode.
3376
3377
slow_mode:  when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3378
3379
    assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3380
    All counting is done with python objects (no overflows or underflows).
3381
    Advances with:  long_cnt += long_step
3382
    Step may be zero -- effectively a slow version of repeat(cnt).
3383
    Either long_cnt or long_step may be a float, Fraction, or Decimal.
3384
*/
3385
3386
/*[clinic input]
3387
@classmethod
3388
itertools.count.__new__
3389
    start as long_cnt: object(c_default="NULL") = 0
3390
    step as long_step: object(c_default="NULL") = 1
3391
Return a count object whose .__next__() method returns consecutive values.
3392
3393
Equivalent to:
3394
    def count(firstval=0, step=1):
3395
        x = firstval
3396
        while 1:
3397
            yield x
3398
            x += step
3399
[clinic start generated code]*/
3400
3401
static PyObject *
3402
itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
3403
                     PyObject *long_step)
3404
/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
3405
4
{
3406
4
    countobject *lz;
3407
4
    int fast_mode;
3408
4
    Py_ssize_t cnt = 0;
3409
4
    long step;
3410
3411
4
    if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3412
4
        (long_step != NULL && !PyNumber_Check(long_step))) {
3413
0
                    PyErr_SetString(PyExc_TypeError, "a number is required");
3414
0
                    return NULL;
3415
0
    }
3416
3417
4
    fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3418
4
                (long_step == NULL || PyLong_Check(long_step));
3419
3420
    /* If not specified, start defaults to 0 */
3421
4
    if (long_cnt != NULL) {
3422
0
        if (fast_mode) {
3423
0
            assert(PyLong_Check(long_cnt));
3424
0
            cnt = PyLong_AsSsize_t(long_cnt);
3425
0
            if (cnt == -1 && PyErr_Occurred()) {
3426
0
                PyErr_Clear();
3427
0
                fast_mode = 0;
3428
0
            }
3429
0
        }
3430
4
    } else {
3431
4
        cnt = 0;
3432
4
        long_cnt = _PyLong_GetZero();
3433
4
    }
3434
4
    Py_INCREF(long_cnt);
3435
3436
    /* If not specified, step defaults to 1 */
3437
4
    if (long_step == NULL) {
3438
4
        long_step = _PyLong_GetOne();
3439
4
    }
3440
4
    Py_INCREF(long_step);
3441
3442
4
    assert(long_cnt != NULL && long_step != NULL);
3443
3444
    /* Fast mode only works when the step is 1 */
3445
4
    if (fast_mode) {
3446
4
        assert(PyLong_Check(long_step));
3447
4
        step = PyLong_AsLong(long_step);
3448
4
        if (step != 1) {
3449
0
            fast_mode = 0;
3450
0
            if (step == -1 && PyErr_Occurred())
3451
0
                PyErr_Clear();
3452
0
        }
3453
4
    }
3454
3455
4
    if (fast_mode)
3456
4
        Py_CLEAR(long_cnt);
3457
0
    else
3458
0
        cnt = PY_SSIZE_T_MAX;
3459
3460
4
    assert((long_cnt == NULL && fast_mode) ||
3461
4
           (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
3462
4
    assert(!fast_mode ||
3463
4
           (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
3464
3465
    /* create countobject structure */
3466
4
    lz = (countobject *)type->tp_alloc(type, 0);
3467
4
    if (lz == NULL) {
3468
0
        Py_XDECREF(long_cnt);
3469
0
        Py_DECREF(long_step);
3470
0
        return NULL;
3471
0
    }
3472
4
    lz->cnt = cnt;
3473
4
    lz->long_cnt = long_cnt;
3474
4
    lz->long_step = long_step;
3475
3476
4
    return (PyObject *)lz;
3477
4
}
3478
3479
static void
3480
count_dealloc(PyObject *op)
3481
0
{
3482
0
    countobject *lz = countobject_CAST(op);
3483
0
    PyTypeObject *tp = Py_TYPE(lz);
3484
0
    PyObject_GC_UnTrack(lz);
3485
0
    Py_XDECREF(lz->long_cnt);
3486
0
    Py_XDECREF(lz->long_step);
3487
0
    tp->tp_free(lz);
3488
0
    Py_DECREF(tp);
3489
0
}
3490
3491
static int
3492
count_traverse(PyObject *op, visitproc visit, void *arg)
3493
4.26k
{
3494
4.26k
    countobject *lz = countobject_CAST(op);
3495
4.26k
    Py_VISIT(Py_TYPE(lz));
3496
4.26k
    Py_VISIT(lz->long_cnt);
3497
4.26k
    Py_VISIT(lz->long_step);
3498
4.26k
    return 0;
3499
4.26k
}
3500
3501
static PyObject *
3502
count_nextlong(countobject *lz)
3503
0
{
3504
0
    PyObject *long_cnt;
3505
0
    PyObject *stepped_up;
3506
3507
0
    long_cnt = lz->long_cnt;
3508
0
    if (long_cnt == NULL) {
3509
        /* Switch to slow_mode */
3510
0
        long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3511
0
        if (long_cnt == NULL)
3512
0
            return NULL;
3513
0
    }
3514
0
    assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
3515
3516
0
    stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3517
0
    if (stepped_up == NULL)
3518
0
        return NULL;
3519
0
    lz->long_cnt = stepped_up;
3520
0
    return long_cnt;
3521
0
}
3522
3523
static PyObject *
3524
count_next(PyObject *op)
3525
0
{
3526
0
    countobject *lz = countobject_CAST(op);
3527
0
#ifndef Py_GIL_DISABLED
3528
0
    if (lz->cnt == PY_SSIZE_T_MAX)
3529
0
        return count_nextlong(lz);
3530
0
    return PyLong_FromSsize_t(lz->cnt++);
3531
#else
3532
    // free-threading version
3533
    // fast mode uses compare-exchange loop
3534
    // slow mode uses a critical section
3535
    PyObject *returned;
3536
    Py_ssize_t cnt;
3537
3538
    cnt = _Py_atomic_load_ssize_relaxed(&lz->cnt);
3539
    for (;;) {
3540
        if (cnt == PY_SSIZE_T_MAX) {
3541
            Py_BEGIN_CRITICAL_SECTION(lz);
3542
            returned = count_nextlong(lz);
3543
            Py_END_CRITICAL_SECTION();
3544
            return returned;
3545
        }
3546
        if (_Py_atomic_compare_exchange_ssize(&lz->cnt, &cnt, cnt + 1)) {
3547
            return PyLong_FromSsize_t(cnt);
3548
        }
3549
    }
3550
#endif
3551
0
}
3552
3553
static PyObject *
3554
count_repr(PyObject *op)
3555
0
{
3556
0
    countobject *lz = countobject_CAST(op);
3557
0
    if (lz->long_cnt == NULL)
3558
0
        return PyUnicode_FromFormat("%s(%zd)",
3559
0
                                    _PyType_Name(Py_TYPE(lz)), lz->cnt);
3560
3561
0
    if (PyLong_Check(lz->long_step)) {
3562
0
        long step = PyLong_AsLong(lz->long_step);
3563
0
        if (step == -1 && PyErr_Occurred()) {
3564
0
            PyErr_Clear();
3565
0
        }
3566
0
        if (step == 1) {
3567
            /* Don't display step when it is an integer equal to 1 */
3568
0
            return PyUnicode_FromFormat("%s(%R)",
3569
0
                                        _PyType_Name(Py_TYPE(lz)),
3570
0
                                        lz->long_cnt);
3571
0
        }
3572
0
    }
3573
0
    return PyUnicode_FromFormat("%s(%R, %R)",
3574
0
                                _PyType_Name(Py_TYPE(lz)),
3575
0
                                lz->long_cnt, lz->long_step);
3576
0
}
3577
3578
static PyType_Slot count_slots[] = {
3579
    {Py_tp_dealloc, count_dealloc},
3580
    {Py_tp_repr, count_repr},
3581
    {Py_tp_getattro, PyObject_GenericGetAttr},
3582
    {Py_tp_doc, (void *)itertools_count__doc__},
3583
    {Py_tp_traverse, count_traverse},
3584
    {Py_tp_iter, PyObject_SelfIter},
3585
    {Py_tp_iternext, count_next},
3586
    {Py_tp_new, itertools_count},
3587
    {Py_tp_free, PyObject_GC_Del},
3588
    {0, NULL},
3589
};
3590
3591
static PyType_Spec count_spec = {
3592
    .name = "itertools.count",
3593
    .basicsize = sizeof(countobject),
3594
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3595
              Py_TPFLAGS_IMMUTABLETYPE),
3596
    .slots = count_slots,
3597
};
3598
3599
3600
/* repeat object ************************************************************/
3601
3602
typedef struct {
3603
    PyObject_HEAD
3604
    PyObject *element;
3605
    Py_ssize_t cnt;
3606
} repeatobject;
3607
3608
0
#define repeatobject_CAST(op)   ((repeatobject *)(op))
3609
3610
static PyObject *
3611
repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3612
0
{
3613
0
    repeatobject *ro;
3614
0
    PyObject *element;
3615
0
    Py_ssize_t cnt = -1, n_args;
3616
0
    static char *kwargs[] = {"object", "times", NULL};
3617
3618
0
    n_args = PyTuple_GET_SIZE(args);
3619
0
    if (kwds != NULL)
3620
0
        n_args += PyDict_GET_SIZE(kwds);
3621
0
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3622
0
                                     &element, &cnt))
3623
0
        return NULL;
3624
    /* Does user supply times argument? */
3625
0
    if (n_args == 2 && cnt < 0)
3626
0
        cnt = 0;
3627
3628
0
    ro = (repeatobject *)type->tp_alloc(type, 0);
3629
0
    if (ro == NULL)
3630
0
        return NULL;
3631
0
    ro->element = Py_NewRef(element);
3632
0
    ro->cnt = cnt;
3633
0
    return (PyObject *)ro;
3634
0
}
3635
3636
static void
3637
repeat_dealloc(PyObject *op)
3638
0
{
3639
0
    repeatobject *ro = repeatobject_CAST(op);
3640
0
    PyTypeObject *tp = Py_TYPE(ro);
3641
0
    PyObject_GC_UnTrack(ro);
3642
0
    Py_XDECREF(ro->element);
3643
0
    tp->tp_free(ro);
3644
0
    Py_DECREF(tp);
3645
0
}
3646
3647
static int
3648
repeat_traverse(PyObject *op, visitproc visit, void *arg)
3649
0
{
3650
0
    repeatobject *ro = repeatobject_CAST(op);
3651
0
    Py_VISIT(Py_TYPE(ro));
3652
0
    Py_VISIT(ro->element);
3653
0
    return 0;
3654
0
}
3655
3656
static PyObject *
3657
repeat_next(PyObject *op)
3658
0
{
3659
0
    repeatobject *ro = repeatobject_CAST(op);
3660
0
    Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(ro->cnt);
3661
0
    if (cnt == 0) {
3662
0
        return NULL;
3663
0
    }
3664
0
    if (cnt > 0) {
3665
0
        cnt--;
3666
0
        FT_ATOMIC_STORE_SSIZE_RELAXED(ro->cnt, cnt);
3667
0
    }
3668
0
    return Py_NewRef(ro->element);
3669
0
}
3670
3671
static PyObject *
3672
repeat_repr(PyObject *op)
3673
0
{
3674
0
    repeatobject *ro = repeatobject_CAST(op);
3675
0
    if (ro->cnt == -1)
3676
0
        return PyUnicode_FromFormat("%s(%R)",
3677
0
                                    _PyType_Name(Py_TYPE(ro)), ro->element);
3678
0
    else
3679
0
        return PyUnicode_FromFormat("%s(%R, %zd)",
3680
0
                                    _PyType_Name(Py_TYPE(ro)), ro->element,
3681
0
                                    ro->cnt);
3682
0
}
3683
3684
static PyObject *
3685
repeat_len(PyObject *op, PyObject *Py_UNUSED(args))
3686
0
{
3687
0
    repeatobject *ro = repeatobject_CAST(op);
3688
0
    if (ro->cnt == -1) {
3689
0
        PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3690
0
        return NULL;
3691
0
    }
3692
0
    return PyLong_FromSize_t(ro->cnt);
3693
0
}
3694
3695
PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
3696
3697
static PyMethodDef repeat_methods[] = {
3698
    {"__length_hint__", repeat_len, METH_NOARGS, length_hint_doc},
3699
    {NULL,              NULL}           /* sentinel */
3700
};
3701
3702
PyDoc_STRVAR(repeat_doc,
3703
"repeat(object [,times]) -> create an iterator which returns the object\n\
3704
for the specified number of times.  If not specified, returns the object\n\
3705
endlessly.");
3706
3707
static PyType_Slot repeat_slots[] = {
3708
    {Py_tp_dealloc, repeat_dealloc},
3709
    {Py_tp_repr, repeat_repr},
3710
    {Py_tp_getattro, PyObject_GenericGetAttr},
3711
    {Py_tp_doc, (void *)repeat_doc},
3712
    {Py_tp_traverse, repeat_traverse},
3713
    {Py_tp_iter, PyObject_SelfIter},
3714
    {Py_tp_iternext, repeat_next},
3715
    {Py_tp_methods, repeat_methods},
3716
    {Py_tp_new, repeat_new},
3717
    {Py_tp_free, PyObject_GC_Del},
3718
    {0, NULL},
3719
};
3720
3721
static PyType_Spec repeat_spec = {
3722
    .name = "itertools.repeat",
3723
    .basicsize = sizeof(repeatobject),
3724
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3725
              Py_TPFLAGS_IMMUTABLETYPE),
3726
    .slots = repeat_slots,
3727
};
3728
3729
3730
/* ziplongest object *********************************************************/
3731
3732
typedef struct {
3733
    PyObject_HEAD
3734
    Py_ssize_t tuplesize;
3735
    Py_ssize_t numactive;
3736
    PyObject *ittuple;                  /* tuple of iterators */
3737
    PyObject *result;
3738
    PyObject *fillvalue;
3739
} ziplongestobject;
3740
3741
0
#define ziplongestobject_CAST(op)   ((ziplongestobject *)(op))
3742
3743
static PyObject *
3744
zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3745
0
{
3746
0
    ziplongestobject *lz;
3747
0
    Py_ssize_t i;
3748
0
    PyObject *ittuple;  /* tuple of iterators */
3749
0
    PyObject *result;
3750
0
    PyObject *fillvalue = Py_None;
3751
0
    Py_ssize_t tuplesize;
3752
3753
0
    if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
3754
0
        fillvalue = NULL;
3755
0
        if (PyDict_GET_SIZE(kwds) == 1) {
3756
0
            fillvalue = PyDict_GetItemWithError(kwds, &_Py_ID(fillvalue));
3757
0
        }
3758
0
        if (fillvalue == NULL) {
3759
0
            if (!PyErr_Occurred()) {
3760
0
                PyErr_SetString(PyExc_TypeError,
3761
0
                    "zip_longest() got an unexpected keyword argument");
3762
0
            }
3763
0
            return NULL;
3764
0
        }
3765
0
    }
3766
3767
    /* args must be a tuple */
3768
0
    assert(PyTuple_Check(args));
3769
0
    tuplesize = PyTuple_GET_SIZE(args);
3770
3771
    /* obtain iterators */
3772
0
    ittuple = PyTuple_New(tuplesize);
3773
0
    if (ittuple == NULL)
3774
0
        return NULL;
3775
0
    for (i=0; i < tuplesize; i++) {
3776
0
        PyObject *item = PyTuple_GET_ITEM(args, i);
3777
0
        PyObject *it = PyObject_GetIter(item);
3778
0
        if (it == NULL) {
3779
0
            Py_DECREF(ittuple);
3780
0
            return NULL;
3781
0
        }
3782
0
        PyTuple_SET_ITEM(ittuple, i, it);
3783
0
    }
3784
3785
    /* create a result holder */
3786
0
    result = PyTuple_New(tuplesize);
3787
0
    if (result == NULL) {
3788
0
        Py_DECREF(ittuple);
3789
0
        return NULL;
3790
0
    }
3791
0
    for (i=0 ; i < tuplesize ; i++) {
3792
0
        Py_INCREF(Py_None);
3793
0
        PyTuple_SET_ITEM(result, i, Py_None);
3794
0
    }
3795
3796
    /* create ziplongestobject structure */
3797
0
    lz = (ziplongestobject *)type->tp_alloc(type, 0);
3798
0
    if (lz == NULL) {
3799
0
        Py_DECREF(ittuple);
3800
0
        Py_DECREF(result);
3801
0
        return NULL;
3802
0
    }
3803
0
    lz->ittuple = ittuple;
3804
0
    lz->tuplesize = tuplesize;
3805
0
    lz->numactive = tuplesize;
3806
0
    lz->result = result;
3807
0
    lz->fillvalue = Py_NewRef(fillvalue);
3808
0
    return (PyObject *)lz;
3809
0
}
3810
3811
static void
3812
zip_longest_dealloc(PyObject *op)
3813
0
{
3814
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3815
0
    PyTypeObject *tp = Py_TYPE(lz);
3816
0
    PyObject_GC_UnTrack(lz);
3817
0
    Py_XDECREF(lz->ittuple);
3818
0
    Py_XDECREF(lz->result);
3819
0
    Py_XDECREF(lz->fillvalue);
3820
0
    tp->tp_free(lz);
3821
0
    Py_DECREF(tp);
3822
0
}
3823
3824
static int
3825
zip_longest_traverse(PyObject *op, visitproc visit, void *arg)
3826
0
{
3827
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3828
0
    Py_VISIT(Py_TYPE(lz));
3829
0
    Py_VISIT(lz->ittuple);
3830
0
    Py_VISIT(lz->result);
3831
0
    Py_VISIT(lz->fillvalue);
3832
0
    return 0;
3833
0
}
3834
3835
static PyObject *
3836
zip_longest_next(PyObject *op)
3837
0
{
3838
0
    ziplongestobject *lz = ziplongestobject_CAST(op);
3839
0
    Py_ssize_t i;
3840
0
    Py_ssize_t tuplesize = lz->tuplesize;
3841
0
    PyObject *result = lz->result;
3842
0
    PyObject *it;
3843
0
    PyObject *item;
3844
0
    PyObject *olditem;
3845
3846
0
    if (tuplesize == 0)
3847
0
        return NULL;
3848
0
    if (lz->numactive == 0)
3849
0
        return NULL;
3850
0
    if (Py_REFCNT(result) == 1) {
3851
0
        Py_INCREF(result);
3852
0
        for (i=0 ; i < tuplesize ; i++) {
3853
0
            it = PyTuple_GET_ITEM(lz->ittuple, i);
3854
0
            if (it == NULL) {
3855
0
                item = Py_NewRef(lz->fillvalue);
3856
0
            } else {
3857
0
                item = PyIter_Next(it);
3858
0
                if (item == NULL) {
3859
0
                    lz->numactive -= 1;
3860
0
                    if (lz->numactive == 0 || PyErr_Occurred()) {
3861
0
                        lz->numactive = 0;
3862
0
                        Py_DECREF(result);
3863
0
                        return NULL;
3864
0
                    } else {
3865
0
                        item = Py_NewRef(lz->fillvalue);
3866
0
                        PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3867
0
                        Py_DECREF(it);
3868
0
                    }
3869
0
                }
3870
0
            }
3871
0
            olditem = PyTuple_GET_ITEM(result, i);
3872
0
            PyTuple_SET_ITEM(result, i, item);
3873
0
            Py_DECREF(olditem);
3874
0
        }
3875
        // bpo-42536: The GC may have untracked this result tuple. Since we're
3876
        // recycling it, make sure it's tracked again:
3877
0
        _PyTuple_Recycle(result);
3878
0
    } else {
3879
0
        result = PyTuple_New(tuplesize);
3880
0
        if (result == NULL)
3881
0
            return NULL;
3882
0
        for (i=0 ; i < tuplesize ; i++) {
3883
0
            it = PyTuple_GET_ITEM(lz->ittuple, i);
3884
0
            if (it == NULL) {
3885
0
                item = Py_NewRef(lz->fillvalue);
3886
0
            } else {
3887
0
                item = PyIter_Next(it);
3888
0
                if (item == NULL) {
3889
0
                    lz->numactive -= 1;
3890
0
                    if (lz->numactive == 0 || PyErr_Occurred()) {
3891
0
                        lz->numactive = 0;
3892
0
                        Py_DECREF(result);
3893
0
                        return NULL;
3894
0
                    } else {
3895
0
                        item = Py_NewRef(lz->fillvalue);
3896
0
                        PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3897
0
                        Py_DECREF(it);
3898
0
                    }
3899
0
                }
3900
0
            }
3901
0
            PyTuple_SET_ITEM(result, i, item);
3902
0
        }
3903
0
    }
3904
0
    return result;
3905
0
}
3906
3907
PyDoc_STRVAR(zip_longest_doc,
3908
"zip_longest(*iterables, fillvalue=None)\n\
3909
--\n\
3910
\n\
3911
Return a zip_longest object whose .__next__() method returns a tuple where\n\
3912
the i-th element comes from the i-th iterable argument.  The .__next__()\n\
3913
method continues until the longest iterable in the argument sequence\n\
3914
is exhausted and then it raises StopIteration.  When the shorter iterables\n\
3915
are exhausted, the fillvalue is substituted in their place.  The fillvalue\n\
3916
defaults to None or can be specified by a keyword argument.\n\
3917
");
3918
3919
static PyType_Slot ziplongest_slots[] = {
3920
    {Py_tp_dealloc, zip_longest_dealloc},
3921
    {Py_tp_getattro, PyObject_GenericGetAttr},
3922
    {Py_tp_doc, (void *)zip_longest_doc},
3923
    {Py_tp_traverse, zip_longest_traverse},
3924
    {Py_tp_iter, PyObject_SelfIter},
3925
    {Py_tp_iternext, zip_longest_next},
3926
    {Py_tp_new, zip_longest_new},
3927
    {Py_tp_free, PyObject_GC_Del},
3928
    {0, NULL},
3929
};
3930
3931
static PyType_Spec ziplongest_spec = {
3932
    .name = "itertools.zip_longest",
3933
    .basicsize = sizeof(ziplongestobject),
3934
    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3935
              Py_TPFLAGS_IMMUTABLETYPE),
3936
    .slots = ziplongest_slots,
3937
};
3938
3939
3940
/* module level code ********************************************************/
3941
3942
PyDoc_STRVAR(module_doc,
3943
"Functional tools for creating and using iterators.\n\
3944
\n\
3945
Infinite iterators:\n\
3946
count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
3947
cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
3948
repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
3949
\n\
3950
Iterators terminating on the shortest input sequence:\n\
3951
accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
3952
batched(p, n) --> [p0, p1, ..., p_n-1], [p_n, p_n+1, ..., p_2n-1], ...\n\
3953
chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
3954
chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
3955
compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3956
dropwhile(predicate, seq) --> seq[n], seq[n+1], starting when predicate fails\n\
3957
groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
3958
filterfalse(predicate, seq) --> elements of seq where predicate(elem) is False\n\
3959
islice(seq, [start,] stop [, step]) --> elements from\n\
3960
       seq[start:stop:step]\n\
3961
pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
3962
starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
3963
tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
3964
takewhile(predicate, seq) --> seq[0], seq[1], until predicate fails\n\
3965
zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
3966
\n\
3967
Combinatoric generators:\n\
3968
product(p, q, ... [repeat=1]) --> cartesian product\n\
3969
permutations(p[, r])\n\
3970
combinations(p, r)\n\
3971
combinations_with_replacement(p, r)\n\
3972
");
3973
3974
static int
3975
itertoolsmodule_traverse(PyObject *mod, visitproc visit, void *arg)
3976
5.26k
{
3977
5.26k
    itertools_state *state = get_module_state(mod);
3978
5.26k
    Py_VISIT(state->accumulate_type);
3979
5.26k
    Py_VISIT(state->batched_type);
3980
5.26k
    Py_VISIT(state->chain_type);
3981
5.26k
    Py_VISIT(state->combinations_type);
3982
5.26k
    Py_VISIT(state->compress_type);
3983
5.26k
    Py_VISIT(state->count_type);
3984
5.26k
    Py_VISIT(state->cwr_type);
3985
5.26k
    Py_VISIT(state->cycle_type);
3986
5.26k
    Py_VISIT(state->dropwhile_type);
3987
5.26k
    Py_VISIT(state->filterfalse_type);
3988
5.26k
    Py_VISIT(state->groupby_type);
3989
5.26k
    Py_VISIT(state->_grouper_type);
3990
5.26k
    Py_VISIT(state->islice_type);
3991
5.26k
    Py_VISIT(state->pairwise_type);
3992
5.26k
    Py_VISIT(state->permutations_type);
3993
5.26k
    Py_VISIT(state->product_type);
3994
5.26k
    Py_VISIT(state->repeat_type);
3995
5.26k
    Py_VISIT(state->starmap_type);
3996
5.26k
    Py_VISIT(state->takewhile_type);
3997
5.26k
    Py_VISIT(state->tee_type);
3998
5.26k
    Py_VISIT(state->teedataobject_type);
3999
5.26k
    Py_VISIT(state->ziplongest_type);
4000
5.26k
    return 0;
4001
5.26k
}
4002
4003
static int
4004
itertoolsmodule_clear(PyObject *mod)
4005
0
{
4006
0
    itertools_state *state = get_module_state(mod);
4007
0
    Py_CLEAR(state->accumulate_type);
4008
0
    Py_CLEAR(state->batched_type);
4009
0
    Py_CLEAR(state->chain_type);
4010
0
    Py_CLEAR(state->combinations_type);
4011
0
    Py_CLEAR(state->compress_type);
4012
0
    Py_CLEAR(state->count_type);
4013
0
    Py_CLEAR(state->cwr_type);
4014
0
    Py_CLEAR(state->cycle_type);
4015
0
    Py_CLEAR(state->dropwhile_type);
4016
0
    Py_CLEAR(state->filterfalse_type);
4017
0
    Py_CLEAR(state->groupby_type);
4018
0
    Py_CLEAR(state->_grouper_type);
4019
0
    Py_CLEAR(state->islice_type);
4020
0
    Py_CLEAR(state->pairwise_type);
4021
0
    Py_CLEAR(state->permutations_type);
4022
0
    Py_CLEAR(state->product_type);
4023
0
    Py_CLEAR(state->repeat_type);
4024
0
    Py_CLEAR(state->starmap_type);
4025
0
    Py_CLEAR(state->takewhile_type);
4026
0
    Py_CLEAR(state->tee_type);
4027
0
    Py_CLEAR(state->teedataobject_type);
4028
0
    Py_CLEAR(state->ziplongest_type);
4029
0
    return 0;
4030
0
}
4031
4032
static void
4033
itertoolsmodule_free(void *mod)
4034
0
{
4035
0
    (void)itertoolsmodule_clear((PyObject *)mod);
4036
0
}
4037
4038
264
#define ADD_TYPE(module, type, spec)                                     \
4039
264
do {                                                                     \
4040
264
    type = (PyTypeObject *)PyType_FromModuleAndSpec(module, spec, NULL); \
4041
264
    if (type == NULL) {                                                  \
4042
0
        return -1;                                                       \
4043
0
    }                                                                    \
4044
264
    if (PyModule_AddType(module, type) < 0) {                            \
4045
0
        return -1;                                                       \
4046
0
    }                                                                    \
4047
264
} while (0)
4048
4049
static int
4050
itertoolsmodule_exec(PyObject *mod)
4051
12
{
4052
12
    itertools_state *state = get_module_state(mod);
4053
12
    ADD_TYPE(mod, state->accumulate_type, &accumulate_spec);
4054
12
    ADD_TYPE(mod, state->batched_type, &batched_spec);
4055
12
    ADD_TYPE(mod, state->chain_type, &chain_spec);
4056
12
    ADD_TYPE(mod, state->combinations_type, &combinations_spec);
4057
12
    ADD_TYPE(mod, state->compress_type, &compress_spec);
4058
12
    ADD_TYPE(mod, state->count_type, &count_spec);
4059
12
    ADD_TYPE(mod, state->cwr_type, &cwr_spec);
4060
12
    ADD_TYPE(mod, state->cycle_type, &cycle_spec);
4061
12
    ADD_TYPE(mod, state->dropwhile_type, &dropwhile_spec);
4062
12
    ADD_TYPE(mod, state->filterfalse_type, &filterfalse_spec);
4063
12
    ADD_TYPE(mod, state->groupby_type, &groupby_spec);
4064
12
    ADD_TYPE(mod, state->_grouper_type, &_grouper_spec);
4065
12
    ADD_TYPE(mod, state->islice_type, &islice_spec);
4066
12
    ADD_TYPE(mod, state->pairwise_type, &pairwise_spec);
4067
12
    ADD_TYPE(mod, state->permutations_type, &permutations_spec);
4068
12
    ADD_TYPE(mod, state->product_type, &product_spec);
4069
12
    ADD_TYPE(mod, state->repeat_type, &repeat_spec);
4070
12
    ADD_TYPE(mod, state->starmap_type, &starmap_spec);
4071
12
    ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
4072
12
    ADD_TYPE(mod, state->tee_type, &tee_spec);
4073
12
    ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
4074
12
    ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
4075
4076
12
    Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
4077
12
    return 0;
4078
12
}
4079
4080
static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4081
    {Py_mod_exec, itertoolsmodule_exec},
4082
    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
4083
    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
4084
    {0, NULL}
4085
};
4086
4087
static PyMethodDef module_methods[] = {
4088
    ITERTOOLS_TEE_METHODDEF
4089
    {NULL, NULL} /* sentinel */
4090
};
4091
4092
4093
static struct PyModuleDef itertoolsmodule = {
4094
    .m_base = PyModuleDef_HEAD_INIT,
4095
    .m_name = "itertools",
4096
    .m_doc = module_doc,
4097
    .m_size = sizeof(itertools_state),
4098
    .m_methods = module_methods,
4099
    .m_slots = itertoolsmodule_slots,
4100
    .m_traverse = itertoolsmodule_traverse,
4101
    .m_clear = itertoolsmodule_clear,
4102
    .m_free = itertoolsmodule_free,
4103
};
4104
4105
PyMODINIT_FUNC
4106
PyInit_itertools(void)
4107
12
{
4108
12
    return PyModuleDef_Init(&itertoolsmodule);
4109
12
}