Coverage Report

Created: 2025-07-11 06:59

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