blob: 643f47b045046cbfe4fa2ef2743350016200f182 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
Raymond Hettingercc061d02020-11-30 20:42:54 -08002
Raymond Hettingerca3788c2015-08-16 14:49:24 -07003#define PY_SSIZE_T_CLEAN
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004#include "Python.h"
Victor Stinner37834132020-10-27 17:12:53 +01005#include "pycore_long.h" // _PyLong_GetZero()
Brandt Bucher226a0122020-12-04 19:45:57 -08006#include "pycore_object.h" // _PyObject_GC_TRACK()
Victor Stinner384621c2020-06-22 17:27:35 +02007#include "pycore_tuple.h" // _PyTuple_ITEMS()
Victor Stinner4a21e572020-04-15 02:35:41 +02008#include <stddef.h> // offsetof()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000010/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +000011 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +000012*/
13
Tal Einat3286ce42018-09-10 21:33:08 +030014/*[clinic input]
15module itertools
16class itertools.groupby "groupbyobject *" "&groupby_type"
17class itertools._grouper "_grouperobject *" "&_grouper_type"
Tal Einatc4bccd32018-09-12 00:49:13 +030018class itertools.teedataobject "teedataobject *" "&teedataobject_type"
19class itertools._tee "teeobject *" "&tee_type"
20class itertools.cycle "cycleobject *" "&cycle_type"
21class itertools.dropwhile "dropwhileobject *" "&dropwhile_type"
22class itertools.takewhile "takewhileobject *" "&takewhile_type"
23class itertools.starmap "starmapobject *" "&starmap_type"
24class itertools.chain "chainobject *" "&chain_type"
25class itertools.combinations "combinationsobject *" "&combinations_type"
26class itertools.combinations_with_replacement "cwr_object *" "&cwr_type"
27class itertools.permutations "permutationsobject *" "&permutations_type"
28class itertools.accumulate "accumulateobject *" "&accumulate_type"
29class itertools.compress "compressobject *" "&compress_type"
30class itertools.filterfalse "filterfalseobject *" "&filterfalse_type"
31class itertools.count "countobject *" "&count_type"
Raymond Hettingercc061d02020-11-30 20:42:54 -080032class itertools.pairwise "pairwiseobject *" "&pairwise_type"
Tal Einat3286ce42018-09-10 21:33:08 +030033[clinic start generated code]*/
Raymond Hettingercc061d02020-11-30 20:42:54 -080034/*[clinic end generated code: output=da39a3ee5e6b4b0d input=6498ed21fbe1bf94]*/
Tal Einat3286ce42018-09-10 21:33:08 +030035
36static PyTypeObject groupby_type;
37static PyTypeObject _grouper_type;
Tal Einatc4bccd32018-09-12 00:49:13 +030038static PyTypeObject teedataobject_type;
39static PyTypeObject tee_type;
40static PyTypeObject cycle_type;
41static PyTypeObject dropwhile_type;
42static PyTypeObject takewhile_type;
43static PyTypeObject starmap_type;
44static PyTypeObject combinations_type;
45static PyTypeObject cwr_type;
46static PyTypeObject permutations_type;
47static PyTypeObject accumulate_type;
48static PyTypeObject compress_type;
49static PyTypeObject filterfalse_type;
50static PyTypeObject count_type;
Raymond Hettingercc061d02020-11-30 20:42:54 -080051static PyTypeObject pairwise_type;
Tal Einatc4bccd32018-09-12 00:49:13 +030052
Tal Einat3286ce42018-09-10 21:33:08 +030053#include "clinic/itertoolsmodule.c.h"
54
Raymond Hettingercc061d02020-11-30 20:42:54 -080055/* pairwise object ***********************************************************/
56
57typedef struct {
58 PyObject_HEAD
59 PyObject *it;
60 PyObject *old;
61} pairwiseobject;
62
63/*[clinic input]
64@classmethod
65itertools.pairwise.__new__ as pairwise_new
66 iterable: object
67 /
68Return an iterator of overlapping pairs taken from the input iterator.
69
70 s -> (s0,s1), (s1,s2), (s2, s3), ...
71
72[clinic start generated code]*/
73
74static PyObject *
75pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
76/*[clinic end generated code: output=9f0267062d384456 input=6e7c3cddb431a8d6]*/
77{
78 PyObject *it;
79 pairwiseobject *po;
80
81 it = PyObject_GetIter(iterable);
82 if (it == NULL) {
83 return NULL;
84 }
85 po = (pairwiseobject *)type->tp_alloc(type, 0);
86 if (po == NULL) {
87 Py_DECREF(it);
88 return NULL;
89 }
90 po->it = it;
91 po->old = NULL;
92 return (PyObject *)po;
93}
94
95static void
96pairwise_dealloc(pairwiseobject *po)
97{
98 PyObject_GC_UnTrack(po);
99 Py_XDECREF(po->it);
100 Py_XDECREF(po->old);
101 Py_TYPE(po)->tp_free(po);
102}
103
104static int
105pairwise_traverse(pairwiseobject *po, visitproc visit, void *arg)
106{
107 Py_VISIT(po->it);
108 Py_VISIT(po->old);
109 return 0;
110}
111
112static PyObject *
113pairwise_next(pairwiseobject *po)
114{
115 PyObject *it = po->it;
116 PyObject *old = po->old;
117 PyObject *new, *result;
118
119 if (it == NULL) {
120 return NULL;
121 }
122 if (old == NULL) {
123 po->old = old = (*Py_TYPE(it)->tp_iternext)(it);
124 if (old == NULL) {
125 Py_CLEAR(po->it);
126 return NULL;
127 }
128 }
129 new = (*Py_TYPE(it)->tp_iternext)(it);
130 if (new == NULL) {
131 Py_CLEAR(po->it);
132 Py_CLEAR(po->old);
133 return NULL;
134 }
135 /* Future optimization: Reuse the result tuple as we do in enumerate() */
136 result = PyTuple_Pack(2, old, new);
137 Py_SETREF(po->old, new);
138 return result;
139}
140
141static PyTypeObject pairwise_type = {
142 PyVarObject_HEAD_INIT(&PyType_Type, 0)
143 "itertools.pairwise", /* tp_name */
144 sizeof(pairwiseobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)pairwise_dealloc, /* tp_dealloc */
148 0, /* tp_vectorcall_offset */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_as_async */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 pairwise_new__doc__, /* tp_doc */
165 (traverseproc)pairwise_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)pairwise_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 PyType_GenericAlloc, /* tp_alloc */
181 pairwise_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
183};
184
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000185
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700186/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000187
188typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyObject_HEAD
190 PyObject *it;
191 PyObject *keyfunc;
192 PyObject *tgtkey;
193 PyObject *currkey;
194 PyObject *currvalue;
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300195 const void *currgrouper; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000196} groupbyobject;
197
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000198static PyObject *_grouper_create(groupbyobject *, PyObject *);
199
Tal Einat3286ce42018-09-10 21:33:08 +0300200/*[clinic input]
201@classmethod
202itertools.groupby.__new__
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000203
Tal Einat3286ce42018-09-10 21:33:08 +0300204 iterable as it: object
205 Elements to divide into groups according to the key function.
206 key as keyfunc: object = None
207 A function for computing the group category for each element.
208 If the key function is not specified or is None, the element itself
209 is used for grouping.
210
211make an iterator that returns consecutive keys and groups from the iterable
212[clinic start generated code]*/
213
214static PyObject *
215itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
216/*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
217{
218 groupbyobject *gbo;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000219
220 gbo = (groupbyobject *)type->tp_alloc(type, 0);
221 if (gbo == NULL)
222 return NULL;
223 gbo->tgtkey = NULL;
224 gbo->currkey = NULL;
225 gbo->currvalue = NULL;
226 gbo->keyfunc = keyfunc;
227 Py_INCREF(keyfunc);
228 gbo->it = PyObject_GetIter(it);
229 if (gbo->it == NULL) {
230 Py_DECREF(gbo);
231 return NULL;
232 }
233 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000234}
235
236static void
237groupby_dealloc(groupbyobject *gbo)
238{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000239 PyObject_GC_UnTrack(gbo);
240 Py_XDECREF(gbo->it);
241 Py_XDECREF(gbo->keyfunc);
242 Py_XDECREF(gbo->tgtkey);
243 Py_XDECREF(gbo->currkey);
244 Py_XDECREF(gbo->currvalue);
245 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000246}
247
248static int
249groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
250{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000251 Py_VISIT(gbo->it);
252 Py_VISIT(gbo->keyfunc);
253 Py_VISIT(gbo->tgtkey);
254 Py_VISIT(gbo->currkey);
255 Py_VISIT(gbo->currvalue);
256 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000257}
258
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300259Py_LOCAL_INLINE(int)
260groupby_step(groupbyobject *gbo)
261{
262 PyObject *newvalue, *newkey, *oldvalue;
263
264 newvalue = PyIter_Next(gbo->it);
265 if (newvalue == NULL)
266 return -1;
267
268 if (gbo->keyfunc == Py_None) {
269 newkey = newvalue;
270 Py_INCREF(newvalue);
271 } else {
Petr Viktorinffd97532020-02-11 17:46:57 +0100272 newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300273 if (newkey == NULL) {
274 Py_DECREF(newvalue);
275 return -1;
276 }
277 }
278
279 oldvalue = gbo->currvalue;
280 gbo->currvalue = newvalue;
281 Py_XSETREF(gbo->currkey, newkey);
282 Py_XDECREF(oldvalue);
283 return 0;
284}
285
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286static PyObject *
287groupby_next(groupbyobject *gbo)
288{
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300289 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000290
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300291 gbo->currgrouper = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000292 /* skip to next iteration group */
293 for (;;) {
294 if (gbo->currkey == NULL)
295 /* pass */;
296 else if (gbo->tgtkey == NULL)
297 break;
298 else {
299 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000300
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700301 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000302 if (rcmp == -1)
303 return NULL;
304 else if (rcmp == 0)
305 break;
306 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000307
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300308 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000309 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000310 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000311 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300312 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000314 grouper = _grouper_create(gbo, gbo->tgtkey);
315 if (grouper == NULL)
316 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000317
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000318 r = PyTuple_Pack(2, gbo->currkey, grouper);
319 Py_DECREF(grouper);
320 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000321}
322
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000323static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530324groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000325{
326 /* reduce as a 'new' call with an optional 'setstate' if groupby
327 * has started
328 */
329 PyObject *value;
330 if (lz->tgtkey && lz->currkey && lz->currvalue)
331 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
332 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
333 else
334 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
335 lz->it, lz->keyfunc);
336
337 return value;
338}
339
340PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
341
342static PyObject *
343groupby_setstate(groupbyobject *lz, PyObject *state)
344{
345 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300346 if (!PyTuple_Check(state)) {
347 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000348 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300349 }
350 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
351 return NULL;
352 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200353 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300354 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200355 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300356 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200357 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300358 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000359 Py_RETURN_NONE;
360}
361
362PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
363
364static PyMethodDef groupby_methods[] = {
365 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
366 reduce_doc},
367 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
368 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700369 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000370};
371
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000372static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000373 PyVarObject_HEAD_INIT(NULL, 0)
374 "itertools.groupby", /* tp_name */
375 sizeof(groupbyobject), /* tp_basicsize */
376 0, /* tp_itemsize */
377 /* methods */
378 (destructor)groupby_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200379 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000380 0, /* tp_getattr */
381 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200382 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000383 0, /* tp_repr */
384 0, /* tp_as_number */
385 0, /* tp_as_sequence */
386 0, /* tp_as_mapping */
387 0, /* tp_hash */
388 0, /* tp_call */
389 0, /* tp_str */
390 PyObject_GenericGetAttr, /* tp_getattro */
391 0, /* tp_setattro */
392 0, /* tp_as_buffer */
393 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
394 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einat3286ce42018-09-10 21:33:08 +0300395 itertools_groupby__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000396 (traverseproc)groupby_traverse, /* tp_traverse */
397 0, /* tp_clear */
398 0, /* tp_richcompare */
399 0, /* tp_weaklistoffset */
400 PyObject_SelfIter, /* tp_iter */
401 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000402 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000403 0, /* tp_members */
404 0, /* tp_getset */
405 0, /* tp_base */
406 0, /* tp_dict */
407 0, /* tp_descr_get */
408 0, /* tp_descr_set */
409 0, /* tp_dictoffset */
410 0, /* tp_init */
411 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300412 itertools_groupby, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000413 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000414};
415
416
417/* _grouper object (internal) ************************************************/
418
419typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000420 PyObject_HEAD
421 PyObject *parent;
422 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000423} _grouperobject;
424
Tal Einat3286ce42018-09-10 21:33:08 +0300425/*[clinic input]
426@classmethod
427itertools._grouper.__new__
428
429 parent: object(subclass_of='&groupby_type')
430 tgtkey: object
431 /
432[clinic start generated code]*/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000433
434static PyObject *
Tal Einat3286ce42018-09-10 21:33:08 +0300435itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
436 PyObject *tgtkey)
437/*[clinic end generated code: output=462efb1cdebb5914 input=dc180d7771fc8c59]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000438{
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000439 return _grouper_create((groupbyobject*) parent, tgtkey);
440}
441
442static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000443_grouper_create(groupbyobject *parent, PyObject *tgtkey)
444{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000445 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000446
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
448 if (igo == NULL)
449 return NULL;
450 igo->parent = (PyObject *)parent;
451 Py_INCREF(parent);
452 igo->tgtkey = tgtkey;
453 Py_INCREF(tgtkey);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300454 parent->currgrouper = igo; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000455
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000456 PyObject_GC_Track(igo);
457 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000458}
459
460static void
461_grouper_dealloc(_grouperobject *igo)
462{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 PyObject_GC_UnTrack(igo);
464 Py_DECREF(igo->parent);
465 Py_DECREF(igo->tgtkey);
466 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000467}
468
469static int
470_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
471{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000472 Py_VISIT(igo->parent);
473 Py_VISIT(igo->tgtkey);
474 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000475}
476
477static PyObject *
478_grouper_next(_grouperobject *igo)
479{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000480 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300481 PyObject *r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000482 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000483
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300484 if (gbo->currgrouper != igo)
485 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000486 if (gbo->currvalue == NULL) {
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300487 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000488 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000489 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000490
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000491 assert(gbo->currkey != NULL);
492 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
493 if (rcmp <= 0)
494 /* got any error or current group is end */
495 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000496
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000497 r = gbo->currvalue;
498 gbo->currvalue = NULL;
499 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000500
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000502}
503
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000504static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530505_grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000506{
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200507 _Py_IDENTIFIER(iter);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300508 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200509 return Py_BuildValue("N(())", _PyEval_GetBuiltinId(&PyId_iter));
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300510 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700511 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000512}
513
514static PyMethodDef _grouper_methods[] = {
515 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
516 reduce_doc},
517 {NULL, NULL} /* sentinel */
518};
519
520
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000521static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000522 PyVarObject_HEAD_INIT(NULL, 0)
523 "itertools._grouper", /* tp_name */
524 sizeof(_grouperobject), /* tp_basicsize */
525 0, /* tp_itemsize */
526 /* methods */
527 (destructor)_grouper_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200528 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000529 0, /* tp_getattr */
530 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200531 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000532 0, /* tp_repr */
533 0, /* tp_as_number */
534 0, /* tp_as_sequence */
535 0, /* tp_as_mapping */
536 0, /* tp_hash */
537 0, /* tp_call */
538 0, /* tp_str */
539 PyObject_GenericGetAttr, /* tp_getattro */
540 0, /* tp_setattro */
541 0, /* tp_as_buffer */
542 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
543 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700544 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000545 0, /* tp_clear */
546 0, /* tp_richcompare */
547 0, /* tp_weaklistoffset */
548 PyObject_SelfIter, /* tp_iter */
549 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000550 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000551 0, /* tp_members */
552 0, /* tp_getset */
553 0, /* tp_base */
554 0, /* tp_dict */
555 0, /* tp_descr_get */
556 0, /* tp_descr_set */
557 0, /* tp_dictoffset */
558 0, /* tp_init */
559 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300560 itertools__grouper, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000561 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000562};
563
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000564
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700565/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000566
Raymond Hettingerad983e72003-11-12 14:32:26 +0000567/* The teedataobject pre-allocates space for LINKCELLS number of objects.
568 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000569 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000570 the other structure members including PyHEAD overhead. The larger the
571 value, the less memory overhead per object and the less time spent
572 allocating/deallocating new links. The smaller the number, the less
573 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000574*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000575#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000576
577typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000578 PyObject_HEAD
579 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700580 int numread; /* 0 <= numread <= LINKCELLS */
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300581 int running;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000582 PyObject *nextlink;
583 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000584} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000585
586typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000587 PyObject_HEAD
588 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700589 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000590 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000591} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000592
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000593static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000594teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000595{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000596 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000598 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
599 if (tdo == NULL)
600 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000601
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300602 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000603 tdo->numread = 0;
604 tdo->nextlink = NULL;
605 Py_INCREF(it);
606 tdo->it = it;
607 PyObject_GC_Track(tdo);
608 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000609}
610
611static PyObject *
612teedataobject_jumplink(teedataobject *tdo)
613{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000614 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000615 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000616 Py_XINCREF(tdo->nextlink);
617 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000618}
619
620static PyObject *
621teedataobject_getitem(teedataobject *tdo, int i)
622{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000623 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000625 assert(i < LINKCELLS);
626 if (i < tdo->numread)
627 value = tdo->values[i];
628 else {
629 /* this is the lead iterator, so fetch more data */
630 assert(i == tdo->numread);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300631 if (tdo->running) {
632 PyErr_SetString(PyExc_RuntimeError,
633 "cannot re-enter the tee iterator");
634 return NULL;
635 }
636 tdo->running = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000637 value = PyIter_Next(tdo->it);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300638 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000639 if (value == NULL)
640 return NULL;
641 tdo->numread++;
642 tdo->values[i] = value;
643 }
644 Py_INCREF(value);
645 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000646}
647
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000648static int
649teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
650{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000651 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700652
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000653 Py_VISIT(tdo->it);
654 for (i = 0; i < tdo->numread; i++)
655 Py_VISIT(tdo->values[i]);
656 Py_VISIT(tdo->nextlink);
657 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000658}
659
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200660static void
661teedataobject_safe_decref(PyObject *obj)
662{
Andy Lesterdffe4c02020-03-04 07:15:20 -0600663 while (obj && Py_IS_TYPE(obj, &teedataobject_type) &&
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200664 Py_REFCNT(obj) == 1) {
665 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
666 ((teedataobject *)obj)->nextlink = NULL;
667 Py_DECREF(obj);
668 obj = nextlink;
669 }
670 Py_XDECREF(obj);
671}
672
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000673static int
674teedataobject_clear(teedataobject *tdo)
675{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000676 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200677 PyObject *tmp;
678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 Py_CLEAR(tdo->it);
680 for (i=0 ; i<tdo->numread ; i++)
681 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200682 tmp = tdo->nextlink;
683 tdo->nextlink = NULL;
684 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000685 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000686}
687
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000688static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000689teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000690{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000691 PyObject_GC_UnTrack(tdo);
692 teedataobject_clear(tdo);
693 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000694}
695
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000696static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530697teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000698{
699 int i;
700 /* create a temporary list of already iterated values */
701 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700702
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000703 if (!values)
704 return NULL;
705 for (i=0 ; i<tdo->numread ; i++) {
706 Py_INCREF(tdo->values[i]);
707 PyList_SET_ITEM(values, i, tdo->values[i]);
708 }
709 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
710 values,
711 tdo->nextlink ? tdo->nextlink : Py_None);
712}
713
Tal Einatc4bccd32018-09-12 00:49:13 +0300714/*[clinic input]
715@classmethod
716itertools.teedataobject.__new__
717 iterable as it: object
718 values: object(subclass_of='&PyList_Type')
719 next: object
720 /
721Data container common to multiple tee objects.
722[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000723
724static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300725itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
726 PyObject *values, PyObject *next)
727/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000728{
729 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000730 Py_ssize_t i, len;
731
732 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000733
734 tdo = (teedataobject *)teedataobject_newinternal(it);
735 if (!tdo)
736 return NULL;
737
738 len = PyList_GET_SIZE(values);
739 if (len > LINKCELLS)
740 goto err;
741 for (i=0; i<len; i++) {
742 tdo->values[i] = PyList_GET_ITEM(values, i);
743 Py_INCREF(tdo->values[i]);
744 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200745 /* len <= LINKCELLS < INT_MAX */
746 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000747
748 if (len == LINKCELLS) {
749 if (next != Py_None) {
Andy Lester55728702020-03-06 16:53:17 -0600750 if (!Py_IS_TYPE(next, &teedataobject_type))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000751 goto err;
752 assert(tdo->nextlink == NULL);
753 Py_INCREF(next);
754 tdo->nextlink = next;
755 }
756 } else {
757 if (next != Py_None)
758 goto err; /* shouldn't have a next if we are not full */
759 }
760 return (PyObject*)tdo;
761
762err:
763 Py_XDECREF(tdo);
764 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
765 return NULL;
766}
767
768static PyMethodDef teedataobject_methods[] = {
769 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
770 reduce_doc},
771 {NULL, NULL} /* sentinel */
772};
773
Raymond Hettingerad983e72003-11-12 14:32:26 +0000774static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700775 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000776 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000777 sizeof(teedataobject), /* tp_basicsize */
778 0, /* tp_itemsize */
779 /* methods */
780 (destructor)teedataobject_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200781 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000782 0, /* tp_getattr */
783 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200784 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000785 0, /* tp_repr */
786 0, /* tp_as_number */
787 0, /* tp_as_sequence */
788 0, /* tp_as_mapping */
789 0, /* tp_hash */
790 0, /* tp_call */
791 0, /* tp_str */
792 PyObject_GenericGetAttr, /* tp_getattro */
793 0, /* tp_setattro */
794 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700795 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300796 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000797 (traverseproc)teedataobject_traverse, /* tp_traverse */
798 (inquiry)teedataobject_clear, /* tp_clear */
799 0, /* tp_richcompare */
800 0, /* tp_weaklistoffset */
801 0, /* tp_iter */
802 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000803 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000804 0, /* tp_members */
805 0, /* tp_getset */
806 0, /* tp_base */
807 0, /* tp_dict */
808 0, /* tp_descr_get */
809 0, /* tp_descr_set */
810 0, /* tp_dictoffset */
811 0, /* tp_init */
812 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300813 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000814 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000815};
816
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000817
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000818static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000819tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000820{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000821 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000822
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000823 if (to->index >= LINKCELLS) {
824 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200825 if (link == NULL)
826 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300827 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000828 to->index = 0;
829 }
830 value = teedataobject_getitem(to->dataobj, to->index);
831 if (value == NULL)
832 return NULL;
833 to->index++;
834 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000835}
836
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000837static int
838tee_traverse(teeobject *to, visitproc visit, void *arg)
839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000840 Py_VISIT((PyObject *)to->dataobj);
841 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000842}
843
Raymond Hettingerad983e72003-11-12 14:32:26 +0000844static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530845tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000846{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000848
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000849 newto = PyObject_GC_New(teeobject, &tee_type);
850 if (newto == NULL)
851 return NULL;
852 Py_INCREF(to->dataobj);
853 newto->dataobj = to->dataobj;
854 newto->index = to->index;
855 newto->weakreflist = NULL;
856 PyObject_GC_Track(newto);
857 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000858}
859
860PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
861
862static PyObject *
863tee_fromiterable(PyObject *iterable)
864{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000865 teeobject *to;
Miss Islington (bot)645e5272021-07-05 15:34:53 -0700866 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000867
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000868 it = PyObject_GetIter(iterable);
869 if (it == NULL)
870 return NULL;
871 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530872 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000873 goto done;
874 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000875
Miss Islington (bot)645e5272021-07-05 15:34:53 -0700876 PyObject *dataobj = teedataobject_newinternal(it);
877 if (!dataobj) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000878 to = NULL;
879 goto done;
880 }
Miss Islington (bot)645e5272021-07-05 15:34:53 -0700881 to = PyObject_GC_New(teeobject, &tee_type);
882 if (to == NULL) {
883 Py_DECREF(dataobj);
884 goto done;
885 }
886 to->dataobj = (teedataobject *)dataobj;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000887 to->index = 0;
888 to->weakreflist = NULL;
889 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000890done:
Miss Islington (bot)645e5272021-07-05 15:34:53 -0700891 Py_DECREF(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000892 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000893}
894
Tal Einatc4bccd32018-09-12 00:49:13 +0300895/*[clinic input]
896@classmethod
897itertools._tee.__new__
898 iterable: object
899 /
900Iterator wrapped to make it copyable.
901[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000902
Tal Einatc4bccd32018-09-12 00:49:13 +0300903static PyObject *
904itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
905/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
906{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000907 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000908}
909
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000910static int
911tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 if (to->weakreflist != NULL)
914 PyObject_ClearWeakRefs((PyObject *) to);
915 Py_CLEAR(to->dataobj);
916 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000917}
918
919static void
920tee_dealloc(teeobject *to)
921{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000922 PyObject_GC_UnTrack(to);
923 tee_clear(to);
924 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000925}
926
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000927static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530928tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000929{
930 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
931}
932
933static PyObject *
934tee_setstate(teeobject *to, PyObject *state)
935{
936 teedataobject *tdo;
937 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300938 if (!PyTuple_Check(state)) {
939 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000940 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300941 }
942 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
943 return NULL;
944 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000945 if (index < 0 || index > LINKCELLS) {
946 PyErr_SetString(PyExc_ValueError, "Index out of range");
947 return NULL;
948 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200949 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300950 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000951 to->index = index;
952 Py_RETURN_NONE;
953}
954
Raymond Hettingerad983e72003-11-12 14:32:26 +0000955static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700956 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
957 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
958 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000959 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000960};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000961
962static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000963 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000964 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000965 sizeof(teeobject), /* tp_basicsize */
966 0, /* tp_itemsize */
967 /* methods */
968 (destructor)tee_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200969 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000970 0, /* tp_getattr */
971 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200972 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000973 0, /* tp_repr */
974 0, /* tp_as_number */
975 0, /* tp_as_sequence */
976 0, /* tp_as_mapping */
977 0, /* tp_hash */
978 0, /* tp_call */
979 0, /* tp_str */
980 0, /* tp_getattro */
981 0, /* tp_setattro */
982 0, /* tp_as_buffer */
983 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300984 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000985 (traverseproc)tee_traverse, /* tp_traverse */
986 (inquiry)tee_clear, /* tp_clear */
987 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700988 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000989 PyObject_SelfIter, /* tp_iter */
990 (iternextfunc)tee_next, /* tp_iternext */
991 tee_methods, /* tp_methods */
992 0, /* tp_members */
993 0, /* tp_getset */
994 0, /* tp_base */
995 0, /* tp_dict */
996 0, /* tp_descr_get */
997 0, /* tp_descr_set */
998 0, /* tp_dictoffset */
999 0, /* tp_init */
1000 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001001 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001002 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001003};
1004
Tal Einatc4bccd32018-09-12 00:49:13 +03001005/*[clinic input]
1006itertools.tee
1007 iterable: object
1008 n: Py_ssize_t = 2
1009 /
1010Returns a tuple of n independent iterators.
1011[clinic start generated code]*/
1012
Raymond Hettingerad983e72003-11-12 14:32:26 +00001013static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001014itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
1015/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +00001016{
Tal Einatc4bccd32018-09-12 00:49:13 +03001017 Py_ssize_t i;
1018 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02001019 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +00001020
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001021 if (n < 0) {
1022 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
1023 return NULL;
1024 }
1025 result = PyTuple_New(n);
1026 if (result == NULL)
1027 return NULL;
1028 if (n == 0)
1029 return result;
1030 it = PyObject_GetIter(iterable);
1031 if (it == NULL) {
1032 Py_DECREF(result);
1033 return NULL;
1034 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001035
Serhiy Storchakaf320be72018-01-25 10:49:40 +02001036 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001037 Py_DECREF(it);
1038 Py_DECREF(result);
1039 return NULL;
1040 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +02001041 if (copyfunc != NULL) {
1042 copyable = it;
1043 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001044 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001045 copyable = tee_fromiterable(it);
1046 Py_DECREF(it);
1047 if (copyable == NULL) {
1048 Py_DECREF(result);
1049 return NULL;
1050 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001051 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
1052 if (copyfunc == NULL) {
1053 Py_DECREF(copyable);
1054 Py_DECREF(result);
1055 return NULL;
1056 }
1057 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02001058
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001059 PyTuple_SET_ITEM(result, 0, copyable);
1060 for (i = 1; i < n; i++) {
1061 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001062 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001063 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001064 Py_DECREF(result);
1065 return NULL;
1066 }
1067 PyTuple_SET_ITEM(result, i, copyable);
1068 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001069 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001070 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +00001071}
1072
Raymond Hettingerad983e72003-11-12 14:32:26 +00001073
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001074/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001075
1076typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001077 PyObject_HEAD
1078 PyObject *it;
1079 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001080 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001081 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001082} cycleobject;
1083
Tal Einatc4bccd32018-09-12 00:49:13 +03001084/*[clinic input]
1085@classmethod
1086itertools.cycle.__new__
1087 iterable: object
1088 /
1089Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
1090[clinic start generated code]*/
1091
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001092static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001093itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
1094/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001095{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001096 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001097 PyObject *saved;
1098 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001100 /* Get iterator. */
1101 it = PyObject_GetIter(iterable);
1102 if (it == NULL)
1103 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001104
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001105 saved = PyList_New(0);
1106 if (saved == NULL) {
1107 Py_DECREF(it);
1108 return NULL;
1109 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001110
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001111 /* create cycleobject structure */
1112 lz = (cycleobject *)type->tp_alloc(type, 0);
1113 if (lz == NULL) {
1114 Py_DECREF(it);
1115 Py_DECREF(saved);
1116 return NULL;
1117 }
1118 lz->it = it;
1119 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001120 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001121 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001123 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001124}
1125
1126static void
1127cycle_dealloc(cycleobject *lz)
1128{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001129 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001130 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001131 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001132 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001133}
1134
1135static int
1136cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1137{
Hai Shi47a23fc2020-06-07 20:05:36 +08001138 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001139 Py_VISIT(lz->saved);
1140 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001141}
1142
1143static PyObject *
1144cycle_next(cycleobject *lz)
1145{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001146 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001147
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001148 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001149 item = PyIter_Next(lz->it);
1150 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001151 if (lz->firstpass)
1152 return item;
1153 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001154 Py_DECREF(item);
1155 return NULL;
1156 }
1157 return item;
1158 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001159 /* Note: StopIteration is already cleared by PyIter_Next() */
1160 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001161 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001162 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001163 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001164 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001165 return NULL;
1166 item = PyList_GET_ITEM(lz->saved, lz->index);
1167 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001168 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001169 lz->index = 0;
1170 Py_INCREF(item);
1171 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001172}
1173
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001174static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301175cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001176{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001177 /* Create a new cycle with the iterator tuple, then set the saved state */
1178 if (lz->it == NULL) {
1179 PyObject *it = PyObject_GetIter(lz->saved);
1180 if (it == NULL)
1181 return NULL;
1182 if (lz->index != 0) {
1183 _Py_IDENTIFIER(__setstate__);
1184 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1185 "n", lz->index);
1186 if (res == NULL) {
1187 Py_DECREF(it);
1188 return NULL;
1189 }
1190 Py_DECREF(res);
1191 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001192 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001193 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001194 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1195 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001196}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001197
1198static PyObject *
1199cycle_setstate(cycleobject *lz, PyObject *state)
1200{
1201 PyObject *saved=NULL;
1202 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001203 if (!PyTuple_Check(state)) {
1204 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001205 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001206 }
1207 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1208 return NULL;
1209 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001210 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001211 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001212 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001213 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001214 Py_RETURN_NONE;
1215}
1216
1217static PyMethodDef cycle_methods[] = {
1218 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1219 reduce_doc},
1220 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1221 setstate_doc},
1222 {NULL, NULL} /* sentinel */
1223};
1224
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001225static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001226 PyVarObject_HEAD_INIT(NULL, 0)
1227 "itertools.cycle", /* tp_name */
1228 sizeof(cycleobject), /* tp_basicsize */
1229 0, /* tp_itemsize */
1230 /* methods */
1231 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001232 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001233 0, /* tp_getattr */
1234 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001235 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001236 0, /* tp_repr */
1237 0, /* tp_as_number */
1238 0, /* tp_as_sequence */
1239 0, /* tp_as_mapping */
1240 0, /* tp_hash */
1241 0, /* tp_call */
1242 0, /* tp_str */
1243 PyObject_GenericGetAttr, /* tp_getattro */
1244 0, /* tp_setattro */
1245 0, /* tp_as_buffer */
1246 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1247 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001248 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001249 (traverseproc)cycle_traverse, /* tp_traverse */
1250 0, /* tp_clear */
1251 0, /* tp_richcompare */
1252 0, /* tp_weaklistoffset */
1253 PyObject_SelfIter, /* tp_iter */
1254 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001255 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001256 0, /* tp_members */
1257 0, /* tp_getset */
1258 0, /* tp_base */
1259 0, /* tp_dict */
1260 0, /* tp_descr_get */
1261 0, /* tp_descr_set */
1262 0, /* tp_dictoffset */
1263 0, /* tp_init */
1264 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001265 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001267};
1268
1269
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001270/* dropwhile object **********************************************************/
1271
1272typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001273 PyObject_HEAD
1274 PyObject *func;
1275 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001276 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001277} dropwhileobject;
1278
Tal Einatc4bccd32018-09-12 00:49:13 +03001279/*[clinic input]
1280@classmethod
1281itertools.dropwhile.__new__
1282 predicate as func: object
1283 iterable as seq: object
1284 /
1285Drop items from the iterable while predicate(item) is true.
1286
1287Afterwards, return every element until the iterable is exhausted.
1288[clinic start generated code]*/
1289
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001290static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001291itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1292/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001293{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001294 PyObject *it;
1295 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001296
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001297 /* Get iterator. */
1298 it = PyObject_GetIter(seq);
1299 if (it == NULL)
1300 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001301
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001302 /* create dropwhileobject structure */
1303 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1304 if (lz == NULL) {
1305 Py_DECREF(it);
1306 return NULL;
1307 }
1308 Py_INCREF(func);
1309 lz->func = func;
1310 lz->it = it;
1311 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001312
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001313 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314}
1315
1316static void
1317dropwhile_dealloc(dropwhileobject *lz)
1318{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001319 PyObject_GC_UnTrack(lz);
1320 Py_XDECREF(lz->func);
1321 Py_XDECREF(lz->it);
1322 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001323}
1324
1325static int
1326dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1327{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001328 Py_VISIT(lz->it);
1329 Py_VISIT(lz->func);
1330 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001331}
1332
1333static PyObject *
1334dropwhile_next(dropwhileobject *lz)
1335{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001336 PyObject *item, *good;
1337 PyObject *it = lz->it;
1338 long ok;
1339 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001341 iternext = *Py_TYPE(it)->tp_iternext;
1342 for (;;) {
1343 item = iternext(it);
1344 if (item == NULL)
1345 return NULL;
1346 if (lz->start == 1)
1347 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001348
Petr Viktorinffd97532020-02-11 17:46:57 +01001349 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001350 if (good == NULL) {
1351 Py_DECREF(item);
1352 return NULL;
1353 }
1354 ok = PyObject_IsTrue(good);
1355 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001356 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001357 lz->start = 1;
1358 return item;
1359 }
1360 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001361 if (ok < 0)
1362 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364}
1365
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001366static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301367dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001368{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001369 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001370}
1371
1372static PyObject *
1373dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1374{
1375 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001376 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001377 return NULL;
1378 lz->start = start;
1379 Py_RETURN_NONE;
1380}
1381
1382static PyMethodDef dropwhile_methods[] = {
1383 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1384 reduce_doc},
1385 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1386 setstate_doc},
1387 {NULL, NULL} /* sentinel */
1388};
1389
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001390static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001391 PyVarObject_HEAD_INIT(NULL, 0)
1392 "itertools.dropwhile", /* tp_name */
1393 sizeof(dropwhileobject), /* tp_basicsize */
1394 0, /* tp_itemsize */
1395 /* methods */
1396 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001397 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001398 0, /* tp_getattr */
1399 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001400 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001401 0, /* tp_repr */
1402 0, /* tp_as_number */
1403 0, /* tp_as_sequence */
1404 0, /* tp_as_mapping */
1405 0, /* tp_hash */
1406 0, /* tp_call */
1407 0, /* tp_str */
1408 PyObject_GenericGetAttr, /* tp_getattro */
1409 0, /* tp_setattro */
1410 0, /* tp_as_buffer */
1411 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1412 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001413 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001414 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001415 0, /* tp_clear */
1416 0, /* tp_richcompare */
1417 0, /* tp_weaklistoffset */
1418 PyObject_SelfIter, /* tp_iter */
1419 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001420 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001421 0, /* tp_members */
1422 0, /* tp_getset */
1423 0, /* tp_base */
1424 0, /* tp_dict */
1425 0, /* tp_descr_get */
1426 0, /* tp_descr_set */
1427 0, /* tp_dictoffset */
1428 0, /* tp_init */
1429 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001430 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001431 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432};
1433
1434
1435/* takewhile object **********************************************************/
1436
1437typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001438 PyObject_HEAD
1439 PyObject *func;
1440 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001441 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442} takewhileobject;
1443
Tal Einatc4bccd32018-09-12 00:49:13 +03001444/*[clinic input]
1445@classmethod
1446itertools.takewhile.__new__
1447 predicate as func: object
1448 iterable as seq: object
1449 /
1450Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1451[clinic start generated code]*/
1452
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001454itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1455/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001457 PyObject *it;
1458 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001460 /* Get iterator. */
1461 it = PyObject_GetIter(seq);
1462 if (it == NULL)
1463 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001465 /* create takewhileobject structure */
1466 lz = (takewhileobject *)type->tp_alloc(type, 0);
1467 if (lz == NULL) {
1468 Py_DECREF(it);
1469 return NULL;
1470 }
1471 Py_INCREF(func);
1472 lz->func = func;
1473 lz->it = it;
1474 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001476 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477}
1478
1479static void
1480takewhile_dealloc(takewhileobject *lz)
1481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 PyObject_GC_UnTrack(lz);
1483 Py_XDECREF(lz->func);
1484 Py_XDECREF(lz->it);
1485 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486}
1487
1488static int
1489takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1490{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001491 Py_VISIT(lz->it);
1492 Py_VISIT(lz->func);
1493 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494}
1495
1496static PyObject *
1497takewhile_next(takewhileobject *lz)
1498{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001499 PyObject *item, *good;
1500 PyObject *it = lz->it;
1501 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001502
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001503 if (lz->stop == 1)
1504 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001506 item = (*Py_TYPE(it)->tp_iternext)(it);
1507 if (item == NULL)
1508 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001509
Petr Viktorinffd97532020-02-11 17:46:57 +01001510 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001511 if (good == NULL) {
1512 Py_DECREF(item);
1513 return NULL;
1514 }
1515 ok = PyObject_IsTrue(good);
1516 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001517 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001518 return item;
1519 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001520 if (ok == 0)
1521 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001522 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523}
1524
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001525static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301526takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001527{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001528 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001529}
1530
1531static PyObject *
1532takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1533{
1534 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001535
Antoine Pitrou721738f2012-08-15 23:20:39 +02001536 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001537 return NULL;
1538 lz->stop = stop;
1539 Py_RETURN_NONE;
1540}
1541
1542static PyMethodDef takewhile_reduce_methods[] = {
1543 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1544 reduce_doc},
1545 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1546 setstate_doc},
1547 {NULL, NULL} /* sentinel */
1548};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001549
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001550static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001551 PyVarObject_HEAD_INIT(NULL, 0)
1552 "itertools.takewhile", /* tp_name */
1553 sizeof(takewhileobject), /* tp_basicsize */
1554 0, /* tp_itemsize */
1555 /* methods */
1556 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001557 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001558 0, /* tp_getattr */
1559 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001560 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001561 0, /* tp_repr */
1562 0, /* tp_as_number */
1563 0, /* tp_as_sequence */
1564 0, /* tp_as_mapping */
1565 0, /* tp_hash */
1566 0, /* tp_call */
1567 0, /* tp_str */
1568 PyObject_GenericGetAttr, /* tp_getattro */
1569 0, /* tp_setattro */
1570 0, /* tp_as_buffer */
1571 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1572 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001573 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001574 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001575 0, /* tp_clear */
1576 0, /* tp_richcompare */
1577 0, /* tp_weaklistoffset */
1578 PyObject_SelfIter, /* tp_iter */
1579 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001580 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001581 0, /* tp_members */
1582 0, /* tp_getset */
1583 0, /* tp_base */
1584 0, /* tp_dict */
1585 0, /* tp_descr_get */
1586 0, /* tp_descr_set */
1587 0, /* tp_dictoffset */
1588 0, /* tp_init */
1589 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001590 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001591 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001592};
1593
1594
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001595/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596
1597typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001598 PyObject_HEAD
1599 PyObject *it;
1600 Py_ssize_t next;
1601 Py_ssize_t stop;
1602 Py_ssize_t step;
1603 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001604} isliceobject;
1605
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001606static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607
1608static PyObject *
1609islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1610{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001611 PyObject *seq;
1612 Py_ssize_t start=0, stop=-1, step=1;
1613 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1614 Py_ssize_t numargs;
1615 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001617 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001618 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001619
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001620 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1621 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001622
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001623 numargs = PyTuple_Size(args);
1624 if (numargs == 2) {
1625 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001626 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001627 if (stop == -1) {
1628 if (PyErr_Occurred())
1629 PyErr_Clear();
1630 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001631 "Stop argument for islice() must be None or "
1632 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001633 return NULL;
1634 }
1635 }
1636 } else {
1637 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001638 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001639 if (start == -1 && PyErr_Occurred())
1640 PyErr_Clear();
1641 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001642 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001643 if (stop == -1) {
1644 if (PyErr_Occurred())
1645 PyErr_Clear();
1646 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001647 "Stop argument for islice() must be None or "
1648 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001649 return NULL;
1650 }
1651 }
1652 }
1653 if (start<0 || stop<-1) {
1654 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001655 "Indices for islice() must be None or "
1656 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001657 return NULL;
1658 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001659
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001660 if (a3 != NULL) {
1661 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001662 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001663 if (step == -1 && PyErr_Occurred())
1664 PyErr_Clear();
1665 }
1666 if (step<1) {
1667 PyErr_SetString(PyExc_ValueError,
1668 "Step for islice() must be a positive integer or None.");
1669 return NULL;
1670 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001672 /* Get iterator. */
1673 it = PyObject_GetIter(seq);
1674 if (it == NULL)
1675 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001677 /* create isliceobject structure */
1678 lz = (isliceobject *)type->tp_alloc(type, 0);
1679 if (lz == NULL) {
1680 Py_DECREF(it);
1681 return NULL;
1682 }
1683 lz->it = it;
1684 lz->next = start;
1685 lz->stop = stop;
1686 lz->step = step;
1687 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001688
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001689 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690}
1691
1692static void
1693islice_dealloc(isliceobject *lz)
1694{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001695 PyObject_GC_UnTrack(lz);
1696 Py_XDECREF(lz->it);
1697 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001698}
1699
1700static int
1701islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1702{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001703 Py_VISIT(lz->it);
1704 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705}
1706
1707static PyObject *
1708islice_next(isliceobject *lz)
1709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001710 PyObject *item;
1711 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001712 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001713 Py_ssize_t oldnext;
1714 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001715
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001716 if (it == NULL)
1717 return NULL;
1718
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001719 iternext = *Py_TYPE(it)->tp_iternext;
1720 while (lz->cnt < lz->next) {
1721 item = iternext(it);
1722 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001723 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001724 Py_DECREF(item);
1725 lz->cnt++;
1726 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001727 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001728 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001729 item = iternext(it);
1730 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001731 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001732 lz->cnt++;
1733 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001734 /* The (size_t) cast below avoids the danger of undefined
1735 behaviour from signed integer overflow. */
1736 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001737 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1738 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001739 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001740
1741empty:
1742 Py_CLEAR(lz->it);
1743 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001744}
1745
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001746static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301747islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001748{
1749 /* When unpickled, generate a new object with the same bounds,
1750 * then 'setstate' with the next and count
1751 */
1752 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001753
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001754 if (lz->it == NULL) {
1755 PyObject *empty_list;
1756 PyObject *empty_it;
1757 empty_list = PyList_New(0);
1758 if (empty_list == NULL)
1759 return NULL;
1760 empty_it = PyObject_GetIter(empty_list);
1761 Py_DECREF(empty_list);
1762 if (empty_it == NULL)
1763 return NULL;
1764 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1765 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001766 if (lz->stop == -1) {
1767 stop = Py_None;
1768 Py_INCREF(stop);
1769 } else {
1770 stop = PyLong_FromSsize_t(lz->stop);
1771 if (stop == NULL)
1772 return NULL;
1773 }
1774 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1775 lz->it, lz->next, stop, lz->step,
1776 lz->cnt);
1777}
1778
1779static PyObject *
1780islice_setstate(isliceobject *lz, PyObject *state)
1781{
1782 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001783
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001784 if (cnt == -1 && PyErr_Occurred())
1785 return NULL;
1786 lz->cnt = cnt;
1787 Py_RETURN_NONE;
1788}
1789
1790static PyMethodDef islice_methods[] = {
1791 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1792 reduce_doc},
1793 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1794 setstate_doc},
1795 {NULL, NULL} /* sentinel */
1796};
1797
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001798PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001799"islice(iterable, stop) --> islice object\n\
1800islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801\n\
1802Return an iterator whose next() method returns selected values from an\n\
1803iterable. If start is specified, will skip all preceding elements;\n\
1804otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001805specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806skipped between successive calls. Works like a slice() on a list\n\
1807but returns an iterator.");
1808
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001809static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001810 PyVarObject_HEAD_INIT(NULL, 0)
1811 "itertools.islice", /* tp_name */
1812 sizeof(isliceobject), /* tp_basicsize */
1813 0, /* tp_itemsize */
1814 /* methods */
1815 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001816 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001817 0, /* tp_getattr */
1818 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001819 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001820 0, /* tp_repr */
1821 0, /* tp_as_number */
1822 0, /* tp_as_sequence */
1823 0, /* tp_as_mapping */
1824 0, /* tp_hash */
1825 0, /* tp_call */
1826 0, /* tp_str */
1827 PyObject_GenericGetAttr, /* tp_getattro */
1828 0, /* tp_setattro */
1829 0, /* tp_as_buffer */
1830 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1831 Py_TPFLAGS_BASETYPE, /* tp_flags */
1832 islice_doc, /* tp_doc */
1833 (traverseproc)islice_traverse, /* tp_traverse */
1834 0, /* tp_clear */
1835 0, /* tp_richcompare */
1836 0, /* tp_weaklistoffset */
1837 PyObject_SelfIter, /* tp_iter */
1838 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001839 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001840 0, /* tp_members */
1841 0, /* tp_getset */
1842 0, /* tp_base */
1843 0, /* tp_dict */
1844 0, /* tp_descr_get */
1845 0, /* tp_descr_set */
1846 0, /* tp_dictoffset */
1847 0, /* tp_init */
1848 0, /* tp_alloc */
1849 islice_new, /* tp_new */
1850 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001851};
1852
1853
1854/* starmap object ************************************************************/
1855
1856typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001857 PyObject_HEAD
1858 PyObject *func;
1859 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001860} starmapobject;
1861
Tal Einatc4bccd32018-09-12 00:49:13 +03001862/*[clinic input]
1863@classmethod
1864itertools.starmap.__new__
1865 function as func: object
1866 iterable as seq: object
1867 /
1868Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1869[clinic start generated code]*/
1870
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001871static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001872itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1873/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001874{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001875 PyObject *it;
1876 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001878 /* Get iterator. */
1879 it = PyObject_GetIter(seq);
1880 if (it == NULL)
1881 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001882
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001883 /* create starmapobject structure */
1884 lz = (starmapobject *)type->tp_alloc(type, 0);
1885 if (lz == NULL) {
1886 Py_DECREF(it);
1887 return NULL;
1888 }
1889 Py_INCREF(func);
1890 lz->func = func;
1891 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001892
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001893 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001894}
1895
1896static void
1897starmap_dealloc(starmapobject *lz)
1898{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001899 PyObject_GC_UnTrack(lz);
1900 Py_XDECREF(lz->func);
1901 Py_XDECREF(lz->it);
1902 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001903}
1904
1905static int
1906starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1907{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001908 Py_VISIT(lz->it);
1909 Py_VISIT(lz->func);
1910 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001911}
1912
1913static PyObject *
1914starmap_next(starmapobject *lz)
1915{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001916 PyObject *args;
1917 PyObject *result;
1918 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001919
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001920 args = (*Py_TYPE(it)->tp_iternext)(it);
1921 if (args == NULL)
1922 return NULL;
1923 if (!PyTuple_CheckExact(args)) {
1924 PyObject *newargs = PySequence_Tuple(args);
1925 Py_DECREF(args);
1926 if (newargs == NULL)
1927 return NULL;
1928 args = newargs;
1929 }
1930 result = PyObject_Call(lz->func, args, NULL);
1931 Py_DECREF(args);
1932 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001933}
1934
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001935static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301936starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001937{
1938 /* Just pickle the iterator */
1939 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1940}
1941
1942static PyMethodDef starmap_methods[] = {
1943 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1944 reduce_doc},
1945 {NULL, NULL} /* sentinel */
1946};
1947
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001948static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001949 PyVarObject_HEAD_INIT(NULL, 0)
1950 "itertools.starmap", /* tp_name */
1951 sizeof(starmapobject), /* tp_basicsize */
1952 0, /* tp_itemsize */
1953 /* methods */
1954 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001955 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001956 0, /* tp_getattr */
1957 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001958 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001959 0, /* tp_repr */
1960 0, /* tp_as_number */
1961 0, /* tp_as_sequence */
1962 0, /* tp_as_mapping */
1963 0, /* tp_hash */
1964 0, /* tp_call */
1965 0, /* tp_str */
1966 PyObject_GenericGetAttr, /* tp_getattro */
1967 0, /* tp_setattro */
1968 0, /* tp_as_buffer */
1969 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1970 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001971 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001972 (traverseproc)starmap_traverse, /* tp_traverse */
1973 0, /* tp_clear */
1974 0, /* tp_richcompare */
1975 0, /* tp_weaklistoffset */
1976 PyObject_SelfIter, /* tp_iter */
1977 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001978 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001979 0, /* tp_members */
1980 0, /* tp_getset */
1981 0, /* tp_base */
1982 0, /* tp_dict */
1983 0, /* tp_descr_get */
1984 0, /* tp_descr_set */
1985 0, /* tp_dictoffset */
1986 0, /* tp_init */
1987 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001988 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001989 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001990};
1991
1992
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001993/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001994
1995typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001996 PyObject_HEAD
1997 PyObject *source; /* Iterator over input iterables */
1998 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001999} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002000
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002001static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002002
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002003static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00002004chain_new_internal(PyTypeObject *type, PyObject *source)
2005{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002006 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00002007
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002008 lz = (chainobject *)type->tp_alloc(type, 0);
2009 if (lz == NULL) {
2010 Py_DECREF(source);
2011 return NULL;
2012 }
2013
2014 lz->source = source;
2015 lz->active = NULL;
2016 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00002017}
2018
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002019static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002020chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002021{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002022 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002023
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03002024 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002025 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002026
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002027 source = PyObject_GetIter(args);
2028 if (source == NULL)
2029 return NULL;
2030
2031 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00002032}
2033
Tal Einatc4bccd32018-09-12 00:49:13 +03002034/*[clinic input]
2035@classmethod
2036itertools.chain.from_iterable
2037 iterable as arg: object
2038 /
2039Alternative chain() constructor taking a single iterable argument that evaluates lazily.
2040[clinic start generated code]*/
2041
Christian Heimesf16baeb2008-02-29 14:57:44 +00002042static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002043itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
2044/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00002045{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002046 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002047
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002048 source = PyObject_GetIter(arg);
2049 if (source == NULL)
2050 return NULL;
2051
2052 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002053}
2054
2055static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002056chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002057{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002058 PyObject_GC_UnTrack(lz);
2059 Py_XDECREF(lz->active);
2060 Py_XDECREF(lz->source);
2061 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002062}
2063
Raymond Hettinger2012f172003-02-07 05:32:58 +00002064static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002065chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00002066{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002067 Py_VISIT(lz->source);
2068 Py_VISIT(lz->active);
2069 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002070}
2071
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002072static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002073chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002074{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002075 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002076
T. Wouters5466d4a2017-03-30 09:58:35 -07002077 /* lz->source is the iterator of iterables. If it's NULL, we've already
2078 * consumed them all. lz->active is the current iterator. If it's NULL,
2079 * we should grab a new one from lz->source. */
2080 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002081 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07002082 PyObject *iterable = PyIter_Next(lz->source);
2083 if (iterable == NULL) {
2084 Py_CLEAR(lz->source);
2085 return NULL; /* no more input sources */
2086 }
2087 lz->active = PyObject_GetIter(iterable);
2088 Py_DECREF(iterable);
2089 if (lz->active == NULL) {
2090 Py_CLEAR(lz->source);
2091 return NULL; /* input not iterable */
2092 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002093 }
T. Wouters5466d4a2017-03-30 09:58:35 -07002094 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
2095 if (item != NULL)
2096 return item;
2097 if (PyErr_Occurred()) {
2098 if (PyErr_ExceptionMatches(PyExc_StopIteration))
2099 PyErr_Clear();
2100 else
2101 return NULL; /* input raised an exception */
2102 }
2103 /* lz->active is consumed, try with the next iterable. */
2104 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002105 }
T. Wouters5466d4a2017-03-30 09:58:35 -07002106 /* Everything had been consumed already. */
2107 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002108}
2109
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002110static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302111chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002112{
2113 if (lz->source) {
2114 /* we can't pickle function objects (itertools.from_iterable) so
2115 * we must use setstate to replace the iterable. One day we
2116 * will fix pickling of functions
2117 */
2118 if (lz->active) {
2119 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
2120 } else {
2121 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
2122 }
2123 } else {
2124 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
2125 }
2126 return NULL;
2127}
2128
2129static PyObject *
2130chain_setstate(chainobject *lz, PyObject *state)
2131{
2132 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002133
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002134 if (!PyTuple_Check(state)) {
2135 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002136 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002137 }
2138 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2139 return NULL;
2140 }
2141 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2142 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2143 return NULL;
2144 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002145
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002146 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002147 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002148 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002149 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002150 Py_RETURN_NONE;
2151}
2152
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002153PyDoc_STRVAR(chain_doc,
2154"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002155\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002156Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002157first iterable until it is exhausted, then elements from the next\n\
2158iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002159
Christian Heimesf16baeb2008-02-29 14:57:44 +00002160static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002161 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002162 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2163 reduce_doc},
2164 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2165 setstate_doc},
Ethan Smitha8403d02020-04-09 20:28:08 -07002166 {"__class_getitem__", (PyCFunction)Py_GenericAlias,
2167 METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002168 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002169};
2170
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002171static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002172 PyVarObject_HEAD_INIT(NULL, 0)
2173 "itertools.chain", /* tp_name */
2174 sizeof(chainobject), /* tp_basicsize */
2175 0, /* tp_itemsize */
2176 /* methods */
2177 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002178 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002179 0, /* tp_getattr */
2180 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002181 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002182 0, /* tp_repr */
2183 0, /* tp_as_number */
2184 0, /* tp_as_sequence */
2185 0, /* tp_as_mapping */
2186 0, /* tp_hash */
2187 0, /* tp_call */
2188 0, /* tp_str */
2189 PyObject_GenericGetAttr, /* tp_getattro */
2190 0, /* tp_setattro */
2191 0, /* tp_as_buffer */
2192 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2193 Py_TPFLAGS_BASETYPE, /* tp_flags */
2194 chain_doc, /* tp_doc */
2195 (traverseproc)chain_traverse, /* tp_traverse */
2196 0, /* tp_clear */
2197 0, /* tp_richcompare */
2198 0, /* tp_weaklistoffset */
2199 PyObject_SelfIter, /* tp_iter */
2200 (iternextfunc)chain_next, /* tp_iternext */
2201 chain_methods, /* tp_methods */
2202 0, /* tp_members */
2203 0, /* tp_getset */
2204 0, /* tp_base */
2205 0, /* tp_dict */
2206 0, /* tp_descr_get */
2207 0, /* tp_descr_set */
2208 0, /* tp_dictoffset */
2209 0, /* tp_init */
2210 0, /* tp_alloc */
2211 chain_new, /* tp_new */
2212 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002213};
2214
2215
Christian Heimesc3f30c42008-02-22 16:37:40 +00002216/* product object ************************************************************/
2217
2218typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002219 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002220 PyObject *pools; /* tuple of pool tuples */
2221 Py_ssize_t *indices; /* one index per pool */
2222 PyObject *result; /* most recently returned result tuple */
2223 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002224} productobject;
2225
2226static PyTypeObject product_type;
2227
2228static PyObject *
2229product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2230{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002231 productobject *lz;
2232 Py_ssize_t nargs, npools, repeat=1;
2233 PyObject *pools = NULL;
2234 Py_ssize_t *indices = NULL;
2235 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002237 if (kwds != NULL) {
2238 char *kwlist[] = {"repeat", 0};
2239 PyObject *tmpargs = PyTuple_New(0);
2240 if (tmpargs == NULL)
2241 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002242 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2243 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002244 Py_DECREF(tmpargs);
2245 return NULL;
2246 }
2247 Py_DECREF(tmpargs);
2248 if (repeat < 0) {
2249 PyErr_SetString(PyExc_ValueError,
2250 "repeat argument cannot be negative");
2251 return NULL;
2252 }
2253 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002254
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002255 assert(PyTuple_CheckExact(args));
2256 if (repeat == 0) {
2257 nargs = 0;
2258 } else {
2259 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002260 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002261 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2262 return NULL;
2263 }
2264 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002265 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002266
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002267 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002268 if (indices == NULL) {
2269 PyErr_NoMemory();
2270 goto error;
2271 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002272
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002273 pools = PyTuple_New(npools);
2274 if (pools == NULL)
2275 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002276
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002277 for (i=0; i < nargs ; ++i) {
2278 PyObject *item = PyTuple_GET_ITEM(args, i);
2279 PyObject *pool = PySequence_Tuple(item);
2280 if (pool == NULL)
2281 goto error;
2282 PyTuple_SET_ITEM(pools, i, pool);
2283 indices[i] = 0;
2284 }
2285 for ( ; i < npools; ++i) {
2286 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2287 Py_INCREF(pool);
2288 PyTuple_SET_ITEM(pools, i, pool);
2289 indices[i] = 0;
2290 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002291
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002292 /* create productobject structure */
2293 lz = (productobject *)type->tp_alloc(type, 0);
2294 if (lz == NULL)
2295 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002296
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002297 lz->pools = pools;
2298 lz->indices = indices;
2299 lz->result = NULL;
2300 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002301
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002302 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002303
2304error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002305 if (indices != NULL)
2306 PyMem_Free(indices);
2307 Py_XDECREF(pools);
2308 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002309}
2310
2311static void
2312product_dealloc(productobject *lz)
2313{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002314 PyObject_GC_UnTrack(lz);
2315 Py_XDECREF(lz->pools);
2316 Py_XDECREF(lz->result);
2317 if (lz->indices != NULL)
2318 PyMem_Free(lz->indices);
2319 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002320}
2321
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002322static PyObject *
2323product_sizeof(productobject *lz, void *unused)
2324{
2325 Py_ssize_t res;
2326
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002327 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002328 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2329 return PyLong_FromSsize_t(res);
2330}
2331
2332PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2333
Christian Heimesc3f30c42008-02-22 16:37:40 +00002334static int
2335product_traverse(productobject *lz, visitproc visit, void *arg)
2336{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002337 Py_VISIT(lz->pools);
2338 Py_VISIT(lz->result);
2339 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002340}
2341
2342static PyObject *
2343product_next(productobject *lz)
2344{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002345 PyObject *pool;
2346 PyObject *elem;
2347 PyObject *oldelem;
2348 PyObject *pools = lz->pools;
2349 PyObject *result = lz->result;
2350 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2351 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002352
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002353 if (lz->stopped)
2354 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002355
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002356 if (result == NULL) {
2357 /* On the first pass, return an initial tuple filled with the
2358 first element from each pool. */
2359 result = PyTuple_New(npools);
2360 if (result == NULL)
2361 goto empty;
2362 lz->result = result;
2363 for (i=0; i < npools; i++) {
2364 pool = PyTuple_GET_ITEM(pools, i);
2365 if (PyTuple_GET_SIZE(pool) == 0)
2366 goto empty;
2367 elem = PyTuple_GET_ITEM(pool, 0);
2368 Py_INCREF(elem);
2369 PyTuple_SET_ITEM(result, i, elem);
2370 }
2371 } else {
2372 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002373
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002374 /* Copy the previous result tuple or re-use it if available */
2375 if (Py_REFCNT(result) > 1) {
2376 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002377 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002378 if (result == NULL)
2379 goto empty;
2380 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002381 Py_DECREF(old_result);
2382 }
Brandt Bucher226a0122020-12-04 19:45:57 -08002383 // bpo-42536: The GC may have untracked this result tuple. Since we're
2384 // recycling it, make sure it's tracked again:
2385 else if (!_PyObject_GC_IS_TRACKED(result)) {
2386 _PyObject_GC_TRACK(result);
2387 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002388 /* Now, we've got the only copy so we can update it in-place */
2389 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002390
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002391 /* Update the pool indices right-to-left. Only advance to the
2392 next pool when the previous one rolls-over */
2393 for (i=npools-1 ; i >= 0 ; i--) {
2394 pool = PyTuple_GET_ITEM(pools, i);
2395 indices[i]++;
2396 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2397 /* Roll-over and advance to next pool */
2398 indices[i] = 0;
2399 elem = PyTuple_GET_ITEM(pool, 0);
2400 Py_INCREF(elem);
2401 oldelem = PyTuple_GET_ITEM(result, i);
2402 PyTuple_SET_ITEM(result, i, elem);
2403 Py_DECREF(oldelem);
2404 } else {
2405 /* No rollover. Just increment and stop here. */
2406 elem = PyTuple_GET_ITEM(pool, indices[i]);
2407 Py_INCREF(elem);
2408 oldelem = PyTuple_GET_ITEM(result, i);
2409 PyTuple_SET_ITEM(result, i, elem);
2410 Py_DECREF(oldelem);
2411 break;
2412 }
2413 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002414
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002415 /* If i is negative, then the indices have all rolled-over
2416 and we're done. */
2417 if (i < 0)
2418 goto empty;
2419 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002420
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002421 Py_INCREF(result);
2422 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002423
2424empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002425 lz->stopped = 1;
2426 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002427}
2428
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002429static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302430product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002431{
2432 if (lz->stopped) {
2433 return Py_BuildValue("O(())", Py_TYPE(lz));
2434 } else if (lz->result == NULL) {
2435 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2436 } else {
2437 PyObject *indices;
2438 Py_ssize_t n, i;
2439
2440 /* we must pickle the indices use them for setstate, and
2441 * additionally indicate that the iterator has started
2442 */
2443 n = PyTuple_GET_SIZE(lz->pools);
2444 indices = PyTuple_New(n);
2445 if (indices == NULL)
2446 return NULL;
2447 for (i=0; i<n; i++){
2448 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2449 if (!index) {
2450 Py_DECREF(indices);
2451 return NULL;
2452 }
2453 PyTuple_SET_ITEM(indices, i, index);
2454 }
2455 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2456 }
2457}
2458
2459static PyObject *
2460product_setstate(productobject *lz, PyObject *state)
2461{
2462 PyObject *result;
2463 Py_ssize_t n, i;
2464
2465 n = PyTuple_GET_SIZE(lz->pools);
2466 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2467 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2468 return NULL;
2469 }
2470 for (i=0; i<n; i++)
2471 {
2472 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2473 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002474 PyObject* pool;
2475 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002476 if (index < 0 && PyErr_Occurred())
2477 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002478 pool = PyTuple_GET_ITEM(lz->pools, i);
2479 poolsize = PyTuple_GET_SIZE(pool);
2480 if (poolsize == 0) {
2481 lz->stopped = 1;
2482 Py_RETURN_NONE;
2483 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002484 /* clamp the index */
2485 if (index < 0)
2486 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002487 else if (index > poolsize-1)
2488 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002489 lz->indices[i] = index;
2490 }
2491
2492 result = PyTuple_New(n);
2493 if (!result)
2494 return NULL;
2495 for (i=0; i<n; i++) {
2496 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2497 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2498 Py_INCREF(element);
2499 PyTuple_SET_ITEM(result, i, element);
2500 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002501 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002502 Py_RETURN_NONE;
2503}
2504
2505static PyMethodDef product_methods[] = {
2506 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2507 reduce_doc},
2508 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2509 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002510 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2511 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002512 {NULL, NULL} /* sentinel */
2513};
2514
Christian Heimesc3f30c42008-02-22 16:37:40 +00002515PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002516"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002517\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002518Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002519For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2520The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2521cycle in a manner similar to an odometer (with the rightmost element changing\n\
2522on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002523To compute the product of an iterable with itself, specify the number\n\
2524of repetitions with the optional repeat keyword argument. For example,\n\
2525product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002526product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2527product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2528
2529static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002530 PyVarObject_HEAD_INIT(NULL, 0)
2531 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002532 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002533 0, /* tp_itemsize */
2534 /* methods */
2535 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002536 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002537 0, /* tp_getattr */
2538 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002539 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002540 0, /* tp_repr */
2541 0, /* tp_as_number */
2542 0, /* tp_as_sequence */
2543 0, /* tp_as_mapping */
2544 0, /* tp_hash */
2545 0, /* tp_call */
2546 0, /* tp_str */
2547 PyObject_GenericGetAttr, /* tp_getattro */
2548 0, /* tp_setattro */
2549 0, /* tp_as_buffer */
2550 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2551 Py_TPFLAGS_BASETYPE, /* tp_flags */
2552 product_doc, /* tp_doc */
2553 (traverseproc)product_traverse, /* tp_traverse */
2554 0, /* tp_clear */
2555 0, /* tp_richcompare */
2556 0, /* tp_weaklistoffset */
2557 PyObject_SelfIter, /* tp_iter */
2558 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002559 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002560 0, /* tp_members */
2561 0, /* tp_getset */
2562 0, /* tp_base */
2563 0, /* tp_dict */
2564 0, /* tp_descr_get */
2565 0, /* tp_descr_set */
2566 0, /* tp_dictoffset */
2567 0, /* tp_init */
2568 0, /* tp_alloc */
2569 product_new, /* tp_new */
2570 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002571};
2572
2573
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002574/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002575
2576typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002577 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002578 PyObject *pool; /* input converted to a tuple */
2579 Py_ssize_t *indices; /* one index per result element */
2580 PyObject *result; /* most recently returned result tuple */
2581 Py_ssize_t r; /* size of result tuple */
2582 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002583} combinationsobject;
2584
Tal Einatc4bccd32018-09-12 00:49:13 +03002585
2586/*[clinic input]
2587@classmethod
2588itertools.combinations.__new__
2589 iterable: object
2590 r: Py_ssize_t
2591Return successive r-length combinations of elements in the iterable.
2592
2593combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2594[clinic start generated code]*/
2595
Christian Heimes380f7f22008-02-28 11:19:05 +00002596static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002597itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2598 Py_ssize_t r)
2599/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002600{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002601 combinationsobject *co;
2602 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002603 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002604 Py_ssize_t *indices = NULL;
2605 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002606
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002607 pool = PySequence_Tuple(iterable);
2608 if (pool == NULL)
2609 goto error;
2610 n = PyTuple_GET_SIZE(pool);
2611 if (r < 0) {
2612 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2613 goto error;
2614 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002615
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002616 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002617 if (indices == NULL) {
2618 PyErr_NoMemory();
2619 goto error;
2620 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002621
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002622 for (i=0 ; i<r ; i++)
2623 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002624
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002625 /* create combinationsobject structure */
2626 co = (combinationsobject *)type->tp_alloc(type, 0);
2627 if (co == NULL)
2628 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002630 co->pool = pool;
2631 co->indices = indices;
2632 co->result = NULL;
2633 co->r = r;
2634 co->stopped = r > n ? 1 : 0;
2635
2636 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002637
2638error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002639 if (indices != NULL)
2640 PyMem_Free(indices);
2641 Py_XDECREF(pool);
2642 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002643}
2644
2645static void
2646combinations_dealloc(combinationsobject *co)
2647{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002648 PyObject_GC_UnTrack(co);
2649 Py_XDECREF(co->pool);
2650 Py_XDECREF(co->result);
2651 if (co->indices != NULL)
2652 PyMem_Free(co->indices);
2653 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002654}
2655
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002656static PyObject *
2657combinations_sizeof(combinationsobject *co, void *unused)
2658{
2659 Py_ssize_t res;
2660
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002661 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002662 res += co->r * sizeof(Py_ssize_t);
2663 return PyLong_FromSsize_t(res);
2664}
2665
Christian Heimes380f7f22008-02-28 11:19:05 +00002666static int
2667combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2668{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002669 Py_VISIT(co->pool);
2670 Py_VISIT(co->result);
2671 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002672}
2673
2674static PyObject *
2675combinations_next(combinationsobject *co)
2676{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002677 PyObject *elem;
2678 PyObject *oldelem;
2679 PyObject *pool = co->pool;
2680 Py_ssize_t *indices = co->indices;
2681 PyObject *result = co->result;
2682 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2683 Py_ssize_t r = co->r;
2684 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002685
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002686 if (co->stopped)
2687 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002688
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002689 if (result == NULL) {
2690 /* On the first pass, initialize result tuple using the indices */
2691 result = PyTuple_New(r);
2692 if (result == NULL)
2693 goto empty;
2694 co->result = result;
2695 for (i=0; i<r ; i++) {
2696 index = indices[i];
2697 elem = PyTuple_GET_ITEM(pool, index);
2698 Py_INCREF(elem);
2699 PyTuple_SET_ITEM(result, i, elem);
2700 }
2701 } else {
2702 /* Copy the previous result tuple or re-use it if available */
2703 if (Py_REFCNT(result) > 1) {
2704 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002705 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002706 if (result == NULL)
2707 goto empty;
2708 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002709 Py_DECREF(old_result);
2710 }
Brandt Bucher226a0122020-12-04 19:45:57 -08002711 // bpo-42536: The GC may have untracked this result tuple. Since we're
2712 // recycling it, make sure it's tracked again:
2713 else if (!_PyObject_GC_IS_TRACKED(result)) {
2714 _PyObject_GC_TRACK(result);
2715 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002716 /* Now, we've got the only copy so we can update it in-place
2717 * CPython's empty tuple is a singleton and cached in
2718 * PyTuple's freelist.
2719 */
2720 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002721
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002722 /* Scan indices right-to-left until finding one that is not
2723 at its maximum (i + n - r). */
2724 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2725 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002726
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002727 /* If i is negative, then the indices are all at
2728 their maximum value and we're done. */
2729 if (i < 0)
2730 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002731
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002732 /* Increment the current index which we know is not at its
2733 maximum. Then move back to the right setting each index
2734 to its lowest possible value (one higher than the index
2735 to its left -- this maintains the sort order invariant). */
2736 indices[i]++;
2737 for (j=i+1 ; j<r ; j++)
2738 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002739
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002740 /* Update the result tuple for the new indices
2741 starting with i, the leftmost index that changed */
2742 for ( ; i<r ; i++) {
2743 index = indices[i];
2744 elem = PyTuple_GET_ITEM(pool, index);
2745 Py_INCREF(elem);
2746 oldelem = PyTuple_GET_ITEM(result, i);
2747 PyTuple_SET_ITEM(result, i, elem);
2748 Py_DECREF(oldelem);
2749 }
2750 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002752 Py_INCREF(result);
2753 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002754
2755empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 co->stopped = 1;
2757 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002758}
2759
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002760static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302761combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002762{
2763 if (lz->result == NULL) {
2764 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2765 } else if (lz->stopped) {
2766 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2767 } else {
2768 PyObject *indices;
2769 Py_ssize_t i;
2770
2771 /* we must pickle the indices and use them for setstate */
2772 indices = PyTuple_New(lz->r);
2773 if (!indices)
2774 return NULL;
2775 for (i=0; i<lz->r; i++)
2776 {
2777 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2778 if (!index) {
2779 Py_DECREF(indices);
2780 return NULL;
2781 }
2782 PyTuple_SET_ITEM(indices, i, index);
2783 }
2784
2785 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2786 }
2787}
2788
2789static PyObject *
2790combinations_setstate(combinationsobject *lz, PyObject *state)
2791{
2792 PyObject *result;
2793 Py_ssize_t i;
2794 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2795
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002796 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002797 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2798 return NULL;
2799 }
2800
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002801 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002802 Py_ssize_t max;
2803 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2804 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002805
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002806 if (index == -1 && PyErr_Occurred())
2807 return NULL; /* not an integer */
2808 max = i + n - lz->r;
2809 /* clamp the index (beware of negative max) */
2810 if (index > max)
2811 index = max;
2812 if (index < 0)
2813 index = 0;
2814 lz->indices[i] = index;
2815 }
2816
2817 result = PyTuple_New(lz->r);
2818 if (result == NULL)
2819 return NULL;
2820 for (i=0; i<lz->r; i++) {
2821 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2822 Py_INCREF(element);
2823 PyTuple_SET_ITEM(result, i, element);
2824 }
2825
Serhiy Storchaka48842712016-04-06 09:45:48 +03002826 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002827 Py_RETURN_NONE;
2828}
2829
2830static PyMethodDef combinations_methods[] = {
2831 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2832 reduce_doc},
2833 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2834 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002835 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2836 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002837 {NULL, NULL} /* sentinel */
2838};
2839
Christian Heimes380f7f22008-02-28 11:19:05 +00002840static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002841 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002842 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002843 sizeof(combinationsobject), /* tp_basicsize */
2844 0, /* tp_itemsize */
2845 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002846 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002847 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002848 0, /* tp_getattr */
2849 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002850 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002851 0, /* tp_repr */
2852 0, /* tp_as_number */
2853 0, /* tp_as_sequence */
2854 0, /* tp_as_mapping */
2855 0, /* tp_hash */
2856 0, /* tp_call */
2857 0, /* tp_str */
2858 PyObject_GenericGetAttr, /* tp_getattro */
2859 0, /* tp_setattro */
2860 0, /* tp_as_buffer */
2861 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2862 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002863 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002864 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002865 0, /* tp_clear */
2866 0, /* tp_richcompare */
2867 0, /* tp_weaklistoffset */
2868 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002869 (iternextfunc)combinations_next, /* tp_iternext */
2870 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002871 0, /* tp_members */
2872 0, /* tp_getset */
2873 0, /* tp_base */
2874 0, /* tp_dict */
2875 0, /* tp_descr_get */
2876 0, /* tp_descr_set */
2877 0, /* tp_dictoffset */
2878 0, /* tp_init */
2879 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002880 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002881 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002882};
2883
2884
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002885/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002886
2887/* Equivalent to:
2888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002889 def combinations_with_replacement(iterable, r):
2890 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2891 # number items returned: (n+r-1)! / r! / (n-1)!
2892 pool = tuple(iterable)
2893 n = len(pool)
2894 indices = [0] * r
2895 yield tuple(pool[i] for i in indices)
2896 while 1:
2897 for i in reversed(range(r)):
2898 if indices[i] != n - 1:
2899 break
2900 else:
2901 return
2902 indices[i:] = [indices[i] + 1] * (r - i)
2903 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002905 def combinations_with_replacement2(iterable, r):
2906 'Alternate version that filters from product()'
2907 pool = tuple(iterable)
2908 n = len(pool)
2909 for indices in product(range(n), repeat=r):
2910 if sorted(indices) == list(indices):
2911 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002912*/
2913typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002914 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002915 PyObject *pool; /* input converted to a tuple */
2916 Py_ssize_t *indices; /* one index per result element */
2917 PyObject *result; /* most recently returned result tuple */
2918 Py_ssize_t r; /* size of result tuple */
2919 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002920} cwrobject;
2921
Tal Einatc4bccd32018-09-12 00:49:13 +03002922/*[clinic input]
2923@classmethod
2924itertools.combinations_with_replacement.__new__
2925 iterable: object
2926 r: Py_ssize_t
2927Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2928
2929combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2930[clinic start generated code]*/
2931
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002932static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002933itertools_combinations_with_replacement_impl(PyTypeObject *type,
2934 PyObject *iterable,
2935 Py_ssize_t r)
2936/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002937{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002938 cwrobject *co;
2939 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002940 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002941 Py_ssize_t *indices = NULL;
2942 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002943
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002944 pool = PySequence_Tuple(iterable);
2945 if (pool == NULL)
2946 goto error;
2947 n = PyTuple_GET_SIZE(pool);
2948 if (r < 0) {
2949 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2950 goto error;
2951 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002952
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002953 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002954 if (indices == NULL) {
2955 PyErr_NoMemory();
2956 goto error;
2957 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002958
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002959 for (i=0 ; i<r ; i++)
2960 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002961
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002962 /* create cwrobject structure */
2963 co = (cwrobject *)type->tp_alloc(type, 0);
2964 if (co == NULL)
2965 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002966
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002967 co->pool = pool;
2968 co->indices = indices;
2969 co->result = NULL;
2970 co->r = r;
2971 co->stopped = !n && r;
2972
2973 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002974
2975error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002976 if (indices != NULL)
2977 PyMem_Free(indices);
2978 Py_XDECREF(pool);
2979 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002980}
2981
2982static void
2983cwr_dealloc(cwrobject *co)
2984{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002985 PyObject_GC_UnTrack(co);
2986 Py_XDECREF(co->pool);
2987 Py_XDECREF(co->result);
2988 if (co->indices != NULL)
2989 PyMem_Free(co->indices);
2990 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002991}
2992
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002993static PyObject *
2994cwr_sizeof(cwrobject *co, void *unused)
2995{
2996 Py_ssize_t res;
2997
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002998 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002999 res += co->r * sizeof(Py_ssize_t);
3000 return PyLong_FromSsize_t(res);
3001}
3002
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003003static int
3004cwr_traverse(cwrobject *co, visitproc visit, void *arg)
3005{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003006 Py_VISIT(co->pool);
3007 Py_VISIT(co->result);
3008 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003009}
3010
3011static PyObject *
3012cwr_next(cwrobject *co)
3013{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003014 PyObject *elem;
3015 PyObject *oldelem;
3016 PyObject *pool = co->pool;
3017 Py_ssize_t *indices = co->indices;
3018 PyObject *result = co->result;
3019 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3020 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05003021 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003022
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003023 if (co->stopped)
3024 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003025
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003026 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05003027 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003028 result = PyTuple_New(r);
3029 if (result == NULL)
3030 goto empty;
3031 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07003032 if (n > 0) {
3033 elem = PyTuple_GET_ITEM(pool, 0);
3034 for (i=0; i<r ; i++) {
3035 assert(indices[i] == 0);
3036 Py_INCREF(elem);
3037 PyTuple_SET_ITEM(result, i, elem);
3038 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003039 }
3040 } else {
3041 /* Copy the previous result tuple or re-use it if available */
3042 if (Py_REFCNT(result) > 1) {
3043 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003044 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003045 if (result == NULL)
3046 goto empty;
3047 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003048 Py_DECREF(old_result);
3049 }
Brandt Bucher226a0122020-12-04 19:45:57 -08003050 // bpo-42536: The GC may have untracked this result tuple. Since we're
3051 // recycling it, make sure it's tracked again:
3052 else if (!_PyObject_GC_IS_TRACKED(result)) {
3053 _PyObject_GC_TRACK(result);
3054 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003055 /* Now, we've got the only copy so we can update it in-place CPython's
3056 empty tuple is a singleton and cached in PyTuple's freelist. */
3057 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003058
Tim Peters9edb1682013-09-03 11:49:31 -05003059 /* Scan indices right-to-left until finding one that is not
3060 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
3062 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003063
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003064 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05003065 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 if (i < 0)
3067 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003068
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003069 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05003070 maximum. Then set all to the right to the same value. */
3071 index = indices[i] + 1;
3072 assert(index < n);
3073 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003074 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05003075 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003076 Py_INCREF(elem);
3077 oldelem = PyTuple_GET_ITEM(result, i);
3078 PyTuple_SET_ITEM(result, i, elem);
3079 Py_DECREF(oldelem);
3080 }
3081 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003082
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003083 Py_INCREF(result);
3084 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003085
3086empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003087 co->stopped = 1;
3088 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003089}
3090
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003091static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303092cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003093{
3094 if (lz->result == NULL) {
3095 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
3096 } else if (lz->stopped) {
3097 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
3098 } else {
3099 PyObject *indices;
3100 Py_ssize_t i;
3101
3102 /* we must pickle the indices and use them for setstate */
3103 indices = PyTuple_New(lz->r);
3104 if (!indices)
3105 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003106 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003107 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
3108 if (!index) {
3109 Py_DECREF(indices);
3110 return NULL;
3111 }
3112 PyTuple_SET_ITEM(indices, i, index);
3113 }
3114
3115 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
3116 }
3117}
3118
3119static PyObject *
3120cwr_setstate(cwrobject *lz, PyObject *state)
3121{
3122 PyObject *result;
3123 Py_ssize_t n, i;
3124
3125 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
3126 {
3127 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3128 return NULL;
3129 }
3130
3131 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003132 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003133 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
3134 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003135
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003136 if (index < 0 && PyErr_Occurred())
3137 return NULL; /* not an integer */
3138 /* clamp the index */
3139 if (index < 0)
3140 index = 0;
3141 else if (index > n-1)
3142 index = n-1;
3143 lz->indices[i] = index;
3144 }
3145 result = PyTuple_New(lz->r);
3146 if (result == NULL)
3147 return NULL;
3148 for (i=0; i<lz->r; i++) {
3149 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3150 Py_INCREF(element);
3151 PyTuple_SET_ITEM(result, i, element);
3152 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003153 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003154 Py_RETURN_NONE;
3155}
3156
3157static PyMethodDef cwr_methods[] = {
3158 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3159 reduce_doc},
3160 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3161 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003162 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3163 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003164 {NULL, NULL} /* sentinel */
3165};
3166
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003167static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003168 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003169 "itertools.combinations_with_replacement", /* tp_name */
3170 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003171 0, /* tp_itemsize */
3172 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003173 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003174 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003175 0, /* tp_getattr */
3176 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003177 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003178 0, /* tp_repr */
3179 0, /* tp_as_number */
3180 0, /* tp_as_sequence */
3181 0, /* tp_as_mapping */
3182 0, /* tp_hash */
3183 0, /* tp_call */
3184 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003185 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003186 0, /* tp_setattro */
3187 0, /* tp_as_buffer */
3188 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003189 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003190 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003191 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003192 0, /* tp_clear */
3193 0, /* tp_richcompare */
3194 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003195 PyObject_SelfIter, /* tp_iter */
3196 (iternextfunc)cwr_next, /* tp_iternext */
3197 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003198 0, /* tp_members */
3199 0, /* tp_getset */
3200 0, /* tp_base */
3201 0, /* tp_dict */
3202 0, /* tp_descr_get */
3203 0, /* tp_descr_set */
3204 0, /* tp_dictoffset */
3205 0, /* tp_init */
3206 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003207 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003208 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003209};
3210
3211
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003212/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003213
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003214def permutations(iterable, r=None):
Sergey0078a0c2019-10-29 08:10:24 +03003215 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3216 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003217 pool = tuple(iterable)
3218 n = len(pool)
3219 r = n if r is None else r
Sergey0078a0c2019-10-29 08:10:24 +03003220 if r > n:
3221 return
3222 indices = list(range(n))
3223 cycles = list(range(n, n-r, -1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003224 yield tuple(pool[i] for i in indices[:r])
3225 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003226 for i in reversed(range(r)):
3227 cycles[i] -= 1
3228 if cycles[i] == 0:
3229 indices[i:] = indices[i+1:] + indices[i:i+1]
3230 cycles[i] = n - i
3231 else:
3232 j = cycles[i]
3233 indices[i], indices[-j] = indices[-j], indices[i]
3234 yield tuple(pool[i] for i in indices[:r])
3235 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003236 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003237 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003238*/
3239
3240typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003241 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003242 PyObject *pool; /* input converted to a tuple */
3243 Py_ssize_t *indices; /* one index per element in the pool */
3244 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3245 PyObject *result; /* most recently returned result tuple */
3246 Py_ssize_t r; /* size of result tuple */
3247 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003248} permutationsobject;
3249
Tal Einatc4bccd32018-09-12 00:49:13 +03003250/*[clinic input]
3251@classmethod
3252itertools.permutations.__new__
3253 iterable: object
3254 r as robj: object = None
3255Return successive r-length permutations of elements in the iterable.
3256
3257permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3258[clinic start generated code]*/
3259
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003260static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003261itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3262 PyObject *robj)
3263/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003264{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003265 permutationsobject *po;
3266 Py_ssize_t n;
3267 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003268 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003269 Py_ssize_t *indices = NULL;
3270 Py_ssize_t *cycles = NULL;
3271 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003272
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003273 pool = PySequence_Tuple(iterable);
3274 if (pool == NULL)
3275 goto error;
3276 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003277
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003278 r = n;
3279 if (robj != Py_None) {
3280 if (!PyLong_Check(robj)) {
3281 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3282 goto error;
3283 }
3284 r = PyLong_AsSsize_t(robj);
3285 if (r == -1 && PyErr_Occurred())
3286 goto error;
3287 }
3288 if (r < 0) {
3289 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3290 goto error;
3291 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003292
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003293 indices = PyMem_New(Py_ssize_t, n);
3294 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003295 if (indices == NULL || cycles == NULL) {
3296 PyErr_NoMemory();
3297 goto error;
3298 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003300 for (i=0 ; i<n ; i++)
3301 indices[i] = i;
3302 for (i=0 ; i<r ; i++)
3303 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003304
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003305 /* create permutationsobject structure */
3306 po = (permutationsobject *)type->tp_alloc(type, 0);
3307 if (po == NULL)
3308 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003309
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003310 po->pool = pool;
3311 po->indices = indices;
3312 po->cycles = cycles;
3313 po->result = NULL;
3314 po->r = r;
3315 po->stopped = r > n ? 1 : 0;
3316
3317 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003318
3319error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003320 if (indices != NULL)
3321 PyMem_Free(indices);
3322 if (cycles != NULL)
3323 PyMem_Free(cycles);
3324 Py_XDECREF(pool);
3325 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003326}
3327
3328static void
3329permutations_dealloc(permutationsobject *po)
3330{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003331 PyObject_GC_UnTrack(po);
3332 Py_XDECREF(po->pool);
3333 Py_XDECREF(po->result);
3334 PyMem_Free(po->indices);
3335 PyMem_Free(po->cycles);
3336 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003337}
3338
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003339static PyObject *
3340permutations_sizeof(permutationsobject *po, void *unused)
3341{
3342 Py_ssize_t res;
3343
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003344 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003345 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3346 res += po->r * sizeof(Py_ssize_t);
3347 return PyLong_FromSsize_t(res);
3348}
3349
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003350static int
3351permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3352{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003353 Py_VISIT(po->pool);
3354 Py_VISIT(po->result);
3355 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003356}
3357
3358static PyObject *
3359permutations_next(permutationsobject *po)
3360{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003361 PyObject *elem;
3362 PyObject *oldelem;
3363 PyObject *pool = po->pool;
3364 Py_ssize_t *indices = po->indices;
3365 Py_ssize_t *cycles = po->cycles;
3366 PyObject *result = po->result;
3367 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3368 Py_ssize_t r = po->r;
3369 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003370
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003371 if (po->stopped)
3372 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003373
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003374 if (result == NULL) {
3375 /* On the first pass, initialize result tuple using the indices */
3376 result = PyTuple_New(r);
3377 if (result == NULL)
3378 goto empty;
3379 po->result = result;
3380 for (i=0; i<r ; i++) {
3381 index = indices[i];
3382 elem = PyTuple_GET_ITEM(pool, index);
3383 Py_INCREF(elem);
3384 PyTuple_SET_ITEM(result, i, elem);
3385 }
3386 } else {
3387 if (n == 0)
3388 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003389
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003390 /* Copy the previous result tuple or re-use it if available */
3391 if (Py_REFCNT(result) > 1) {
3392 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003393 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003394 if (result == NULL)
3395 goto empty;
3396 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003397 Py_DECREF(old_result);
3398 }
Brandt Bucher226a0122020-12-04 19:45:57 -08003399 // bpo-42536: The GC may have untracked this result tuple. Since we're
3400 // recycling it, make sure it's tracked again:
3401 else if (!_PyObject_GC_IS_TRACKED(result)) {
3402 _PyObject_GC_TRACK(result);
3403 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003404 /* Now, we've got the only copy so we can update it in-place */
3405 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003406
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003407 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3408 for (i=r-1 ; i>=0 ; i--) {
3409 cycles[i] -= 1;
3410 if (cycles[i] == 0) {
3411 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3412 index = indices[i];
3413 for (j=i ; j<n-1 ; j++)
3414 indices[j] = indices[j+1];
3415 indices[n-1] = index;
3416 cycles[i] = n - i;
3417 } else {
3418 j = cycles[i];
3419 index = indices[i];
3420 indices[i] = indices[n-j];
3421 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003423 for (k=i; k<r ; k++) {
3424 /* start with i, the leftmost element that changed */
3425 /* yield tuple(pool[k] for k in indices[:r]) */
3426 index = indices[k];
3427 elem = PyTuple_GET_ITEM(pool, index);
3428 Py_INCREF(elem);
3429 oldelem = PyTuple_GET_ITEM(result, k);
3430 PyTuple_SET_ITEM(result, k, elem);
3431 Py_DECREF(oldelem);
3432 }
3433 break;
3434 }
3435 }
3436 /* If i is negative, then the cycles have all
3437 rolled-over and we're done. */
3438 if (i < 0)
3439 goto empty;
3440 }
3441 Py_INCREF(result);
3442 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003443
3444empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003445 po->stopped = 1;
3446 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003447}
3448
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003449static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303450permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003451{
3452 if (po->result == NULL) {
3453 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3454 } else if (po->stopped) {
3455 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3456 } else {
3457 PyObject *indices=NULL, *cycles=NULL;
3458 Py_ssize_t n, i;
3459
3460 /* we must pickle the indices and cycles and use them for setstate */
3461 n = PyTuple_GET_SIZE(po->pool);
3462 indices = PyTuple_New(n);
3463 if (indices == NULL)
3464 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003465 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003466 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3467 if (!index)
3468 goto err;
3469 PyTuple_SET_ITEM(indices, i, index);
3470 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003471
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003472 cycles = PyTuple_New(po->r);
3473 if (cycles == NULL)
3474 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003475 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003476 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3477 if (!index)
3478 goto err;
3479 PyTuple_SET_ITEM(cycles, i, index);
3480 }
3481 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3482 po->pool, po->r,
3483 indices, cycles);
3484 err:
3485 Py_XDECREF(indices);
3486 Py_XDECREF(cycles);
3487 return NULL;
3488 }
3489}
3490
3491static PyObject *
3492permutations_setstate(permutationsobject *po, PyObject *state)
3493{
3494 PyObject *indices, *cycles, *result;
3495 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003496
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003497 if (!PyTuple_Check(state)) {
3498 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3499 return NULL;
3500 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003501 if (!PyArg_ParseTuple(state, "O!O!",
3502 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003503 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003504 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003505 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003506
3507 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003508 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003509 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3510 return NULL;
3511 }
3512
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003513 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003514 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3515 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3516 if (index < 0 && PyErr_Occurred())
3517 return NULL; /* not an integer */
3518 /* clamp the index */
3519 if (index < 0)
3520 index = 0;
3521 else if (index > n-1)
3522 index = n-1;
3523 po->indices[i] = index;
3524 }
3525
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003526 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003527 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3528 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3529 if (index < 0 && PyErr_Occurred())
3530 return NULL; /* not an integer */
3531 if (index < 1)
3532 index = 1;
3533 else if (index > n-i)
3534 index = n-i;
3535 po->cycles[i] = index;
3536 }
3537 result = PyTuple_New(po->r);
3538 if (result == NULL)
3539 return NULL;
3540 for (i=0; i<po->r; i++) {
3541 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3542 Py_INCREF(element);
3543 PyTuple_SET_ITEM(result, i, element);
3544 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003545 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003546 Py_RETURN_NONE;
3547}
3548
3549static PyMethodDef permuations_methods[] = {
3550 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3551 reduce_doc},
3552 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3553 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003554 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3555 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003556 {NULL, NULL} /* sentinel */
3557};
3558
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003559static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003560 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003561 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003562 sizeof(permutationsobject), /* tp_basicsize */
3563 0, /* tp_itemsize */
3564 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003565 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003566 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003567 0, /* tp_getattr */
3568 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003569 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003570 0, /* tp_repr */
3571 0, /* tp_as_number */
3572 0, /* tp_as_sequence */
3573 0, /* tp_as_mapping */
3574 0, /* tp_hash */
3575 0, /* tp_call */
3576 0, /* tp_str */
3577 PyObject_GenericGetAttr, /* tp_getattro */
3578 0, /* tp_setattro */
3579 0, /* tp_as_buffer */
3580 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3581 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003582 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003583 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003584 0, /* tp_clear */
3585 0, /* tp_richcompare */
3586 0, /* tp_weaklistoffset */
3587 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003588 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003589 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003590 0, /* tp_members */
3591 0, /* tp_getset */
3592 0, /* tp_base */
3593 0, /* tp_dict */
3594 0, /* tp_descr_get */
3595 0, /* tp_descr_set */
3596 0, /* tp_dictoffset */
3597 0, /* tp_init */
3598 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003599 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003600 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003601};
3602
Tal Einatc4bccd32018-09-12 00:49:13 +03003603
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003604/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003605
3606typedef struct {
3607 PyObject_HEAD
3608 PyObject *total;
3609 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003610 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003611 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003612} accumulateobject;
3613
Tal Einatc4bccd32018-09-12 00:49:13 +03003614/*[clinic input]
3615@classmethod
3616itertools.accumulate.__new__
3617 iterable: object
3618 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003619 *
3620 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003621Return series of accumulated sums (or other binary function results).
3622[clinic start generated code]*/
3623
Raymond Hettinger482ba772010-12-01 22:48:00 +00003624static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003625itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003626 PyObject *binop, PyObject *initial)
3627/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003628{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003629 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003630 accumulateobject *lz;
3631
Raymond Hettinger482ba772010-12-01 22:48:00 +00003632 /* Get iterator. */
3633 it = PyObject_GetIter(iterable);
3634 if (it == NULL)
3635 return NULL;
3636
Raymond Hettinger482ba772010-12-01 22:48:00 +00003637 /* create accumulateobject structure */
3638 lz = (accumulateobject *)type->tp_alloc(type, 0);
3639 if (lz == NULL) {
3640 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003641 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003642 }
3643
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003644 if (binop != Py_None) {
3645 Py_XINCREF(binop);
3646 lz->binop = binop;
3647 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003648 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003649 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003650 Py_XINCREF(initial);
3651 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003652 return (PyObject *)lz;
3653}
3654
3655static void
3656accumulate_dealloc(accumulateobject *lz)
3657{
3658 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003659 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003660 Py_XDECREF(lz->total);
3661 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003662 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003663 Py_TYPE(lz)->tp_free(lz);
3664}
3665
3666static int
3667accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3668{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003669 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003670 Py_VISIT(lz->it);
3671 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003672 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003673 return 0;
3674}
3675
3676static PyObject *
3677accumulate_next(accumulateobject *lz)
3678{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003679 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003680
Lisa Roach9718b592018-09-23 17:34:59 -07003681 if (lz->initial != Py_None) {
3682 lz->total = lz->initial;
3683 Py_INCREF(Py_None);
3684 lz->initial = Py_None;
3685 Py_INCREF(lz->total);
3686 return lz->total;
3687 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003688 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003689 if (val == NULL)
3690 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003691
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003692 if (lz->total == NULL) {
3693 Py_INCREF(val);
3694 lz->total = val;
3695 return lz->total;
3696 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003697
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003698 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003699 newtotal = PyNumber_Add(lz->total, val);
3700 else
3701 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003702 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003703 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003704 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003705
Raymond Hettinger482ba772010-12-01 22:48:00 +00003706 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003707 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003708 return newtotal;
3709}
3710
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003711static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303712accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003713{
Lisa Roach9718b592018-09-23 17:34:59 -07003714 if (lz->initial != Py_None) {
3715 PyObject *it;
3716
3717 assert(lz->total == NULL);
3718 if (PyType_Ready(&chain_type) < 0)
3719 return NULL;
3720 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3721 lz->initial, lz->it);
3722 if (it == NULL)
3723 return NULL;
3724 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3725 it, lz->binop?lz->binop:Py_None, Py_None);
3726 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003727 if (lz->total == Py_None) {
3728 PyObject *it;
3729
3730 if (PyType_Ready(&chain_type) < 0)
3731 return NULL;
3732 if (PyType_Ready(&islice_type) < 0)
3733 return NULL;
3734 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3735 lz->total, lz->it);
3736 if (it == NULL)
3737 return NULL;
3738 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3739 it, lz->binop ? lz->binop : Py_None);
3740 if (it == NULL)
3741 return NULL;
3742 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3743 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003744 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3745 lz->it, lz->binop?lz->binop:Py_None,
3746 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003747}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003748
3749static PyObject *
3750accumulate_setstate(accumulateobject *lz, PyObject *state)
3751{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003752 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003753 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003754 Py_RETURN_NONE;
3755}
3756
3757static PyMethodDef accumulate_methods[] = {
3758 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3759 reduce_doc},
3760 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3761 setstate_doc},
3762 {NULL, NULL} /* sentinel */
3763};
3764
Raymond Hettinger482ba772010-12-01 22:48:00 +00003765static PyTypeObject accumulate_type = {
3766 PyVarObject_HEAD_INIT(NULL, 0)
3767 "itertools.accumulate", /* tp_name */
3768 sizeof(accumulateobject), /* tp_basicsize */
3769 0, /* tp_itemsize */
3770 /* methods */
3771 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003772 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003773 0, /* tp_getattr */
3774 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003775 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003776 0, /* tp_repr */
3777 0, /* tp_as_number */
3778 0, /* tp_as_sequence */
3779 0, /* tp_as_mapping */
3780 0, /* tp_hash */
3781 0, /* tp_call */
3782 0, /* tp_str */
3783 PyObject_GenericGetAttr, /* tp_getattro */
3784 0, /* tp_setattro */
3785 0, /* tp_as_buffer */
3786 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3787 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003788 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003789 (traverseproc)accumulate_traverse, /* tp_traverse */
3790 0, /* tp_clear */
3791 0, /* tp_richcompare */
3792 0, /* tp_weaklistoffset */
3793 PyObject_SelfIter, /* tp_iter */
3794 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003795 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003796 0, /* tp_members */
3797 0, /* tp_getset */
3798 0, /* tp_base */
3799 0, /* tp_dict */
3800 0, /* tp_descr_get */
3801 0, /* tp_descr_set */
3802 0, /* tp_dictoffset */
3803 0, /* tp_init */
3804 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003805 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003806 PyObject_GC_Del, /* tp_free */
3807};
3808
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003809
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003810/* compress object ************************************************************/
3811
3812/* Equivalent to:
3813
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003814 def compress(data, selectors):
3815 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3816 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003817*/
3818
3819typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003820 PyObject_HEAD
3821 PyObject *data;
3822 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003823} compressobject;
3824
Tal Einatc4bccd32018-09-12 00:49:13 +03003825/*[clinic input]
3826@classmethod
3827itertools.compress.__new__
3828 data as seq1: object
3829 selectors as seq2: object
3830Return data elements corresponding to true selector elements.
3831
3832Forms a shorter iterator from selected data elements using the selectors to
3833choose the data elements.
3834[clinic start generated code]*/
3835
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003836static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003837itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3838/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003840 PyObject *data=NULL, *selectors=NULL;
3841 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003842
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003843 data = PyObject_GetIter(seq1);
3844 if (data == NULL)
3845 goto fail;
3846 selectors = PyObject_GetIter(seq2);
3847 if (selectors == NULL)
3848 goto fail;
3849
3850 /* create compressobject structure */
3851 lz = (compressobject *)type->tp_alloc(type, 0);
3852 if (lz == NULL)
3853 goto fail;
3854 lz->data = data;
3855 lz->selectors = selectors;
3856 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003857
3858fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003859 Py_XDECREF(data);
3860 Py_XDECREF(selectors);
3861 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003862}
3863
3864static void
3865compress_dealloc(compressobject *lz)
3866{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003867 PyObject_GC_UnTrack(lz);
3868 Py_XDECREF(lz->data);
3869 Py_XDECREF(lz->selectors);
3870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003871}
3872
3873static int
3874compress_traverse(compressobject *lz, visitproc visit, void *arg)
3875{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003876 Py_VISIT(lz->data);
3877 Py_VISIT(lz->selectors);
3878 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003879}
3880
3881static PyObject *
3882compress_next(compressobject *lz)
3883{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003884 PyObject *data = lz->data, *selectors = lz->selectors;
3885 PyObject *datum, *selector;
3886 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3887 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3888 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003889
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003890 while (1) {
3891 /* Steps: get datum, get selector, evaluate selector.
3892 Order is important (to match the pure python version
3893 in terms of which input gets a chance to raise an
3894 exception first).
3895 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003896
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003897 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003898 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003899 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003900
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003901 selector = selectornext(selectors);
3902 if (selector == NULL) {
3903 Py_DECREF(datum);
3904 return NULL;
3905 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003907 ok = PyObject_IsTrue(selector);
3908 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003909 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003910 return datum;
3911 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003912 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 return NULL;
3914 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003915}
3916
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003917static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303918compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003919{
3920 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3921 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003922}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003923
3924static PyMethodDef compress_methods[] = {
3925 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3926 reduce_doc},
3927 {NULL, NULL} /* sentinel */
3928};
3929
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003930static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003931 PyVarObject_HEAD_INIT(NULL, 0)
3932 "itertools.compress", /* tp_name */
3933 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003934 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 /* methods */
3936 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003937 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003938 0, /* tp_getattr */
3939 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003940 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003941 0, /* tp_repr */
3942 0, /* tp_as_number */
3943 0, /* tp_as_sequence */
3944 0, /* tp_as_mapping */
3945 0, /* tp_hash */
3946 0, /* tp_call */
3947 0, /* tp_str */
3948 PyObject_GenericGetAttr, /* tp_getattro */
3949 0, /* tp_setattro */
3950 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003951 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003952 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003953 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003954 (traverseproc)compress_traverse, /* tp_traverse */
3955 0, /* tp_clear */
3956 0, /* tp_richcompare */
3957 0, /* tp_weaklistoffset */
3958 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003959 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003960 compress_methods, /* tp_methods */
3961 0, /* tp_members */
3962 0, /* tp_getset */
3963 0, /* tp_base */
3964 0, /* tp_dict */
3965 0, /* tp_descr_get */
3966 0, /* tp_descr_set */
3967 0, /* tp_dictoffset */
3968 0, /* tp_init */
3969 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003970 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003971 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003972};
3973
3974
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003975/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003976
3977typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003978 PyObject_HEAD
3979 PyObject *func;
3980 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003981} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003982
Tal Einatc4bccd32018-09-12 00:49:13 +03003983/*[clinic input]
3984@classmethod
3985itertools.filterfalse.__new__
3986 function as func: object
3987 iterable as seq: object
3988 /
3989Return those items of iterable for which function(item) is false.
3990
3991If function is None, return the items that are false.
3992[clinic start generated code]*/
3993
Raymond Hettinger60eca932003-02-09 06:40:58 +00003994static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003995itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3996/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003997{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003998 PyObject *it;
3999 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004001 /* Get iterator. */
4002 it = PyObject_GetIter(seq);
4003 if (it == NULL)
4004 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004006 /* create filterfalseobject structure */
4007 lz = (filterfalseobject *)type->tp_alloc(type, 0);
4008 if (lz == NULL) {
4009 Py_DECREF(it);
4010 return NULL;
4011 }
4012 Py_INCREF(func);
4013 lz->func = func;
4014 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004015
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004016 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004017}
4018
4019static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004020filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00004021{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004022 PyObject_GC_UnTrack(lz);
4023 Py_XDECREF(lz->func);
4024 Py_XDECREF(lz->it);
4025 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00004026}
4027
4028static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004029filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00004030{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004031 Py_VISIT(lz->it);
4032 Py_VISIT(lz->func);
4033 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004034}
4035
4036static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004037filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00004038{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004039 PyObject *item;
4040 PyObject *it = lz->it;
4041 long ok;
4042 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00004043
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004044 iternext = *Py_TYPE(it)->tp_iternext;
4045 for (;;) {
4046 item = iternext(it);
4047 if (item == NULL)
4048 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004049
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004050 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
4051 ok = PyObject_IsTrue(item);
4052 } else {
4053 PyObject *good;
Petr Viktorinffd97532020-02-11 17:46:57 +01004054 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004055 if (good == NULL) {
4056 Py_DECREF(item);
4057 return NULL;
4058 }
4059 ok = PyObject_IsTrue(good);
4060 Py_DECREF(good);
4061 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02004062 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004063 return item;
4064 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02004065 if (ok < 0)
4066 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004067 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00004068}
4069
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004070static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304071filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004072{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004073 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
4074}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004075
4076static PyMethodDef filterfalse_methods[] = {
4077 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
4078 reduce_doc},
4079 {NULL, NULL} /* sentinel */
4080};
4081
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004082static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004083 PyVarObject_HEAD_INIT(NULL, 0)
4084 "itertools.filterfalse", /* tp_name */
4085 sizeof(filterfalseobject), /* tp_basicsize */
4086 0, /* tp_itemsize */
4087 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004088 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004089 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004090 0, /* tp_getattr */
4091 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004092 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004093 0, /* tp_repr */
4094 0, /* tp_as_number */
4095 0, /* tp_as_sequence */
4096 0, /* tp_as_mapping */
4097 0, /* tp_hash */
4098 0, /* tp_call */
4099 0, /* tp_str */
4100 PyObject_GenericGetAttr, /* tp_getattro */
4101 0, /* tp_setattro */
4102 0, /* tp_as_buffer */
4103 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4104 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004105 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004106 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004107 0, /* tp_clear */
4108 0, /* tp_richcompare */
4109 0, /* tp_weaklistoffset */
4110 PyObject_SelfIter, /* tp_iter */
4111 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004112 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004113 0, /* tp_members */
4114 0, /* tp_getset */
4115 0, /* tp_base */
4116 0, /* tp_dict */
4117 0, /* tp_descr_get */
4118 0, /* tp_descr_set */
4119 0, /* tp_dictoffset */
4120 0, /* tp_init */
4121 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004122 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004124};
4125
4126
4127/* count object ************************************************************/
4128
4129typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004130 PyObject_HEAD
4131 Py_ssize_t cnt;
4132 PyObject *long_cnt;
4133 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004134} countobject;
4135
Raymond Hettinger30729212009-02-12 06:28:27 +00004136/* Counting logic and invariants:
4137
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004138fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00004139
Serhiy Storchaka483405b2015-02-17 10:14:30 +02004140 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004141 Advances with: cnt += 1
4142 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00004143
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004144slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004146 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4147 All counting is done with python objects (no overflows or underflows).
4148 Advances with: long_cnt += long_step
4149 Step may be zero -- effectively a slow version of repeat(cnt).
4150 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004151*/
4152
Tal Einatc4bccd32018-09-12 00:49:13 +03004153/*[clinic input]
4154@classmethod
4155itertools.count.__new__
4156 start as long_cnt: object(c_default="NULL") = 0
4157 step as long_step: object(c_default="NULL") = 1
4158Return a count object whose .__next__() method returns consecutive values.
4159
4160Equivalent to:
4161 def count(firstval=0, step=1):
4162 x = firstval
4163 while 1:
4164 yield x
4165 x += step
4166[clinic start generated code]*/
4167
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004168static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004169itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4170 PyObject *long_step)
4171/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004172{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004173 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004174 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004175 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004176 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004178 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4179 (long_step != NULL && !PyNumber_Check(long_step))) {
4180 PyErr_SetString(PyExc_TypeError, "a number is required");
4181 return NULL;
4182 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004183
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004184 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4185 (long_step == NULL || PyLong_Check(long_step));
4186
4187 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004188 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004189 if (fast_mode) {
4190 assert(PyLong_Check(long_cnt));
4191 cnt = PyLong_AsSsize_t(long_cnt);
4192 if (cnt == -1 && PyErr_Occurred()) {
4193 PyErr_Clear();
4194 fast_mode = 0;
4195 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004197 } else {
4198 cnt = 0;
Victor Stinner37834132020-10-27 17:12:53 +01004199 long_cnt = _PyLong_GetZero();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004200 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004201 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004202
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004203 /* If not specified, step defaults to 1 */
Victor Stinner37834132020-10-27 17:12:53 +01004204 if (long_step == NULL) {
4205 long_step = _PyLong_GetOne();
4206 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004207 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004208
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004209 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004210
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004211 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004212 if (fast_mode) {
4213 assert(PyLong_Check(long_step));
4214 step = PyLong_AsLong(long_step);
4215 if (step != 1) {
4216 fast_mode = 0;
4217 if (step == -1 && PyErr_Occurred())
4218 PyErr_Clear();
4219 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004220 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004221
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004222 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004223 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004224 else
4225 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004226
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004227 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4228 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4229 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004230 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004231
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004232 /* create countobject structure */
4233 lz = (countobject *)type->tp_alloc(type, 0);
4234 if (lz == NULL) {
4235 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004236 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004237 return NULL;
4238 }
4239 lz->cnt = cnt;
4240 lz->long_cnt = long_cnt;
4241 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004242
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004243 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004244}
4245
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004246static void
4247count_dealloc(countobject *lz)
4248{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004249 PyObject_GC_UnTrack(lz);
4250 Py_XDECREF(lz->long_cnt);
4251 Py_XDECREF(lz->long_step);
4252 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004253}
4254
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004255static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004256count_traverse(countobject *lz, visitproc visit, void *arg)
4257{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004258 Py_VISIT(lz->long_cnt);
4259 Py_VISIT(lz->long_step);
4260 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004261}
4262
4263static PyObject *
4264count_nextlong(countobject *lz)
4265{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004266 PyObject *long_cnt;
4267 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004268
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004269 long_cnt = lz->long_cnt;
4270 if (long_cnt == NULL) {
4271 /* Switch to slow_mode */
4272 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4273 if (long_cnt == NULL)
4274 return NULL;
4275 }
4276 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004277
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004278 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4279 if (stepped_up == NULL)
4280 return NULL;
4281 lz->long_cnt = stepped_up;
4282 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004283}
4284
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004285static PyObject *
4286count_next(countobject *lz)
4287{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004288 if (lz->cnt == PY_SSIZE_T_MAX)
4289 return count_nextlong(lz);
4290 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004291}
4292
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004293static PyObject *
4294count_repr(countobject *lz)
4295{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004296 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004297 return PyUnicode_FromFormat("%s(%zd)",
4298 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004300 if (PyLong_Check(lz->long_step)) {
4301 long step = PyLong_AsLong(lz->long_step);
4302 if (step == -1 && PyErr_Occurred()) {
4303 PyErr_Clear();
4304 }
4305 if (step == 1) {
4306 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004307 return PyUnicode_FromFormat("%s(%R)",
4308 _PyType_Name(Py_TYPE(lz)),
4309 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004310 }
4311 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004312 return PyUnicode_FromFormat("%s(%R, %R)",
4313 _PyType_Name(Py_TYPE(lz)),
4314 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004315}
4316
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004317static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304318count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004319{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004320 if (lz->cnt == PY_SSIZE_T_MAX)
4321 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4322 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004323}
4324
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004325static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004326 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004327 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004328 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004329};
4330
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004331static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004332 PyVarObject_HEAD_INIT(NULL, 0)
4333 "itertools.count", /* tp_name */
4334 sizeof(countobject), /* tp_basicsize */
4335 0, /* tp_itemsize */
4336 /* methods */
4337 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004338 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004339 0, /* tp_getattr */
4340 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004341 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004342 (reprfunc)count_repr, /* tp_repr */
4343 0, /* tp_as_number */
4344 0, /* tp_as_sequence */
4345 0, /* tp_as_mapping */
4346 0, /* tp_hash */
4347 0, /* tp_call */
4348 0, /* tp_str */
4349 PyObject_GenericGetAttr, /* tp_getattro */
4350 0, /* tp_setattro */
4351 0, /* tp_as_buffer */
4352 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004353 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004354 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004355 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004356 0, /* tp_clear */
4357 0, /* tp_richcompare */
4358 0, /* tp_weaklistoffset */
4359 PyObject_SelfIter, /* tp_iter */
4360 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004361 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004362 0, /* tp_members */
4363 0, /* tp_getset */
4364 0, /* tp_base */
4365 0, /* tp_dict */
4366 0, /* tp_descr_get */
4367 0, /* tp_descr_set */
4368 0, /* tp_dictoffset */
4369 0, /* tp_init */
4370 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004371 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004372 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004373};
4374
4375
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004376/* repeat object ************************************************************/
4377
4378typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004379 PyObject_HEAD
4380 PyObject *element;
4381 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004382} repeatobject;
4383
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004384static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004385
4386static PyObject *
4387repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4388{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004389 repeatobject *ro;
4390 PyObject *element;
Raymond Hettingeraf464502019-11-09 20:28:31 -08004391 Py_ssize_t cnt = -1, n_args;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004392 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004393
Raymond Hettingeraf464502019-11-09 20:28:31 -08004394 n_args = PyTuple_GET_SIZE(args);
4395 if (kwds != NULL)
4396 n_args += PyDict_GET_SIZE(kwds);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004397 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4398 &element, &cnt))
4399 return NULL;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004400 /* Does user supply times argument? */
Raymond Hettingeraf464502019-11-09 20:28:31 -08004401 if (n_args == 2 && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004402 cnt = 0;
4403
4404 ro = (repeatobject *)type->tp_alloc(type, 0);
4405 if (ro == NULL)
4406 return NULL;
4407 Py_INCREF(element);
4408 ro->element = element;
4409 ro->cnt = cnt;
4410 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004411}
4412
4413static void
4414repeat_dealloc(repeatobject *ro)
4415{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004416 PyObject_GC_UnTrack(ro);
4417 Py_XDECREF(ro->element);
4418 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004419}
4420
4421static int
4422repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4423{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004424 Py_VISIT(ro->element);
4425 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004426}
4427
4428static PyObject *
4429repeat_next(repeatobject *ro)
4430{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004431 if (ro->cnt == 0)
4432 return NULL;
4433 if (ro->cnt > 0)
4434 ro->cnt--;
4435 Py_INCREF(ro->element);
4436 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004437}
4438
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004439static PyObject *
4440repeat_repr(repeatobject *ro)
4441{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004442 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004443 return PyUnicode_FromFormat("%s(%R)",
4444 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004445 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004446 return PyUnicode_FromFormat("%s(%R, %zd)",
4447 _PyType_Name(Py_TYPE(ro)), ro->element,
4448 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004449}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004450
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004451static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304452repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004453{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004454 if (ro->cnt == -1) {
4455 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4456 return NULL;
4457 }
4458 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004459}
4460
Armin Rigof5b3e362006-02-11 21:32:43 +00004461PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004462
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004463static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304464repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004465{
4466 /* unpickle this so that a new repeat iterator is constructed with an
4467 * object, then call __setstate__ on it to set cnt
4468 */
4469 if (ro->cnt >= 0)
4470 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4471 else
4472 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4473}
4474
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004475static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004476 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004477 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004478 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004479};
4480
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004481PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004482"repeat(object [,times]) -> create an iterator which returns the object\n\
4483for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004484endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004485
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004486static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004487 PyVarObject_HEAD_INIT(NULL, 0)
4488 "itertools.repeat", /* tp_name */
4489 sizeof(repeatobject), /* tp_basicsize */
4490 0, /* tp_itemsize */
4491 /* methods */
4492 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004493 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004494 0, /* tp_getattr */
4495 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004496 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004497 (reprfunc)repeat_repr, /* tp_repr */
4498 0, /* tp_as_number */
4499 0, /* tp_as_sequence */
4500 0, /* tp_as_mapping */
4501 0, /* tp_hash */
4502 0, /* tp_call */
4503 0, /* tp_str */
4504 PyObject_GenericGetAttr, /* tp_getattro */
4505 0, /* tp_setattro */
4506 0, /* tp_as_buffer */
4507 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4508 Py_TPFLAGS_BASETYPE, /* tp_flags */
4509 repeat_doc, /* tp_doc */
4510 (traverseproc)repeat_traverse, /* tp_traverse */
4511 0, /* tp_clear */
4512 0, /* tp_richcompare */
4513 0, /* tp_weaklistoffset */
4514 PyObject_SelfIter, /* tp_iter */
4515 (iternextfunc)repeat_next, /* tp_iternext */
4516 repeat_methods, /* tp_methods */
4517 0, /* tp_members */
4518 0, /* tp_getset */
4519 0, /* tp_base */
4520 0, /* tp_dict */
4521 0, /* tp_descr_get */
4522 0, /* tp_descr_set */
4523 0, /* tp_dictoffset */
4524 0, /* tp_init */
4525 0, /* tp_alloc */
4526 repeat_new, /* tp_new */
4527 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004528};
4529
Tal Einatc4bccd32018-09-12 00:49:13 +03004530
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004531/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004532
4533typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004534 PyObject_HEAD
4535 Py_ssize_t tuplesize;
4536 Py_ssize_t numactive;
4537 PyObject *ittuple; /* tuple of iterators */
4538 PyObject *result;
4539 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004540} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004541
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004542static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004543
4544static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004545zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004546{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004547 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004548 ziplongestobject *lz;
4549 Py_ssize_t i;
4550 PyObject *ittuple; /* tuple of iterators */
4551 PyObject *result;
4552 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004553 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004554
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004555 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004556 fillvalue = NULL;
4557 if (PyDict_GET_SIZE(kwds) == 1) {
4558 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4559 }
4560 if (fillvalue == NULL) {
4561 if (!PyErr_Occurred()) {
4562 PyErr_SetString(PyExc_TypeError,
4563 "zip_longest() got an unexpected keyword argument");
4564 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004565 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004566 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004567 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004568
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004569 /* args must be a tuple */
4570 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004571 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004572
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004573 /* obtain iterators */
4574 ittuple = PyTuple_New(tuplesize);
4575 if (ittuple == NULL)
4576 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004577 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004578 PyObject *item = PyTuple_GET_ITEM(args, i);
4579 PyObject *it = PyObject_GetIter(item);
4580 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004581 Py_DECREF(ittuple);
4582 return NULL;
4583 }
4584 PyTuple_SET_ITEM(ittuple, i, it);
4585 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004586
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004587 /* create a result holder */
4588 result = PyTuple_New(tuplesize);
4589 if (result == NULL) {
4590 Py_DECREF(ittuple);
4591 return NULL;
4592 }
4593 for (i=0 ; i < tuplesize ; i++) {
4594 Py_INCREF(Py_None);
4595 PyTuple_SET_ITEM(result, i, Py_None);
4596 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004598 /* create ziplongestobject structure */
4599 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4600 if (lz == NULL) {
4601 Py_DECREF(ittuple);
4602 Py_DECREF(result);
4603 return NULL;
4604 }
4605 lz->ittuple = ittuple;
4606 lz->tuplesize = tuplesize;
4607 lz->numactive = tuplesize;
4608 lz->result = result;
4609 Py_INCREF(fillvalue);
4610 lz->fillvalue = fillvalue;
4611 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004612}
4613
4614static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004615zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004616{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004617 PyObject_GC_UnTrack(lz);
4618 Py_XDECREF(lz->ittuple);
4619 Py_XDECREF(lz->result);
4620 Py_XDECREF(lz->fillvalue);
4621 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004622}
4623
4624static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004625zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004626{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004627 Py_VISIT(lz->ittuple);
4628 Py_VISIT(lz->result);
4629 Py_VISIT(lz->fillvalue);
4630 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004631}
4632
4633static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004634zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004635{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004636 Py_ssize_t i;
4637 Py_ssize_t tuplesize = lz->tuplesize;
4638 PyObject *result = lz->result;
4639 PyObject *it;
4640 PyObject *item;
4641 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004643 if (tuplesize == 0)
4644 return NULL;
4645 if (lz->numactive == 0)
4646 return NULL;
4647 if (Py_REFCNT(result) == 1) {
4648 Py_INCREF(result);
4649 for (i=0 ; i < tuplesize ; i++) {
4650 it = PyTuple_GET_ITEM(lz->ittuple, i);
4651 if (it == NULL) {
4652 Py_INCREF(lz->fillvalue);
4653 item = lz->fillvalue;
4654 } else {
4655 item = PyIter_Next(it);
4656 if (item == NULL) {
4657 lz->numactive -= 1;
4658 if (lz->numactive == 0 || PyErr_Occurred()) {
4659 lz->numactive = 0;
4660 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004661 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004662 } else {
4663 Py_INCREF(lz->fillvalue);
4664 item = lz->fillvalue;
4665 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4666 Py_DECREF(it);
4667 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004668 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004669 }
4670 olditem = PyTuple_GET_ITEM(result, i);
4671 PyTuple_SET_ITEM(result, i, item);
4672 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004673 }
Brandt Bucher226a0122020-12-04 19:45:57 -08004674 // bpo-42536: The GC may have untracked this result tuple. Since we're
4675 // recycling it, make sure it's tracked again:
4676 if (!_PyObject_GC_IS_TRACKED(result)) {
4677 _PyObject_GC_TRACK(result);
4678 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004679 } else {
4680 result = PyTuple_New(tuplesize);
4681 if (result == NULL)
4682 return NULL;
4683 for (i=0 ; i < tuplesize ; i++) {
4684 it = PyTuple_GET_ITEM(lz->ittuple, i);
4685 if (it == NULL) {
4686 Py_INCREF(lz->fillvalue);
4687 item = lz->fillvalue;
4688 } else {
4689 item = PyIter_Next(it);
4690 if (item == NULL) {
4691 lz->numactive -= 1;
4692 if (lz->numactive == 0 || PyErr_Occurred()) {
4693 lz->numactive = 0;
4694 Py_DECREF(result);
4695 return NULL;
4696 } else {
4697 Py_INCREF(lz->fillvalue);
4698 item = lz->fillvalue;
4699 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4700 Py_DECREF(it);
4701 }
4702 }
4703 }
4704 PyTuple_SET_ITEM(result, i, item);
4705 }
4706 }
4707 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004708}
4709
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004710static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304711zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004712{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004713
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004714 /* Create a new tuple with empty sequences where appropriate to pickle.
4715 * Then use setstate to set the fillvalue
4716 */
4717 int i;
4718 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004719
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004720 if (args == NULL)
4721 return NULL;
4722 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4723 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4724 if (elem == NULL) {
4725 elem = PyTuple_New(0);
4726 if (elem == NULL) {
4727 Py_DECREF(args);
4728 return NULL;
4729 }
4730 } else
4731 Py_INCREF(elem);
4732 PyTuple_SET_ITEM(args, i, elem);
4733 }
4734 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4735}
4736
4737static PyObject *
4738zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4739{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004740 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004741 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004742 Py_RETURN_NONE;
4743}
4744
4745static PyMethodDef zip_longest_methods[] = {
4746 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4747 reduce_doc},
4748 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4749 setstate_doc},
4750 {NULL, NULL} /* sentinel */
4751};
4752
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004753PyDoc_STRVAR(zip_longest_doc,
4754"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004755\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004756Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004757the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004758method continues until the longest iterable in the argument sequence\n\
4759is exhausted and then it raises StopIteration. When the shorter iterables\n\
4760are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4761defaults to None or can be specified by a keyword argument.\n\
4762");
4763
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004764static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004765 PyVarObject_HEAD_INIT(NULL, 0)
4766 "itertools.zip_longest", /* tp_name */
4767 sizeof(ziplongestobject), /* tp_basicsize */
4768 0, /* tp_itemsize */
4769 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004770 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004771 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004772 0, /* tp_getattr */
4773 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004774 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004775 0, /* tp_repr */
4776 0, /* tp_as_number */
4777 0, /* tp_as_sequence */
4778 0, /* tp_as_mapping */
4779 0, /* tp_hash */
4780 0, /* tp_call */
4781 0, /* tp_str */
4782 PyObject_GenericGetAttr, /* tp_getattro */
4783 0, /* tp_setattro */
4784 0, /* tp_as_buffer */
4785 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4786 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004787 zip_longest_doc, /* tp_doc */
4788 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004789 0, /* tp_clear */
4790 0, /* tp_richcompare */
4791 0, /* tp_weaklistoffset */
4792 PyObject_SelfIter, /* tp_iter */
4793 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004794 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004795 0, /* tp_members */
4796 0, /* tp_getset */
4797 0, /* tp_base */
4798 0, /* tp_dict */
4799 0, /* tp_descr_get */
4800 0, /* tp_descr_set */
4801 0, /* tp_dictoffset */
4802 0, /* tp_init */
4803 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004804 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004805 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004806};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004807
Tal Einatc4bccd32018-09-12 00:49:13 +03004808
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004809/* module level code ********************************************************/
4810
4811PyDoc_STRVAR(module_doc,
4812"Functional tools for creating and using iterators.\n\
4813\n\
4814Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004815count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004816cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004817repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004818\n\
4819Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004820accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004821chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4822chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004823compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4824dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4825groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004826filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004827islice(seq, [start,] stop [, step]) --> elements from\n\
4828 seq[start:stop:step]\n\
Raymond Hettingercc061d02020-11-30 20:42:54 -08004829pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004830starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004831tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004832takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004833zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004834\n\
4835Combinatoric generators:\n\
4836product(p, q, ... [repeat=1]) --> cartesian product\n\
4837permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004838combinations(p, r)\n\
4839combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004840");
4841
Dong-hee Na514c4692020-03-18 02:46:24 +09004842static int
4843itertoolsmodule_exec(PyObject *m)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004844{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004845 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004846 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004847 &combinations_type,
4848 &cwr_type,
4849 &cycle_type,
4850 &dropwhile_type,
4851 &takewhile_type,
4852 &islice_type,
4853 &starmap_type,
4854 &chain_type,
4855 &compress_type,
4856 &filterfalse_type,
4857 &count_type,
4858 &ziplongest_type,
Raymond Hettingercc061d02020-11-30 20:42:54 -08004859 &pairwise_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004860 &permutations_type,
4861 &product_type,
4862 &repeat_type,
4863 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004864 &_grouper_type,
4865 &tee_type,
Dong-hee Na05e4a292020-03-23 01:17:34 +09004866 &teedataobject_type
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004867 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004868
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004869 Py_SET_TYPE(&teedataobject_type, &PyType_Type);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004870
Dong-hee Na05e4a292020-03-23 01:17:34 +09004871 for (size_t i = 0; i < Py_ARRAY_LENGTH(typelist); i++) {
4872 if (PyModule_AddType(m, typelist[i]) < 0) {
Dong-hee Na514c4692020-03-18 02:46:24 +09004873 return -1;
4874 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004875 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004876
Dong-hee Na514c4692020-03-18 02:46:24 +09004877 return 0;
4878}
4879
4880static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4881 {Py_mod_exec, itertoolsmodule_exec},
4882 {0, NULL}
4883};
4884
4885static PyMethodDef module_methods[] = {
4886 ITERTOOLS_TEE_METHODDEF
4887 {NULL, NULL} /* sentinel */
4888};
4889
4890
4891static struct PyModuleDef itertoolsmodule = {
4892 PyModuleDef_HEAD_INIT,
4893 "itertools",
4894 module_doc,
4895 0,
4896 module_methods,
4897 itertoolsmodule_slots,
4898 NULL,
4899 NULL,
4900 NULL
4901};
4902
4903PyMODINIT_FUNC
4904PyInit_itertools(void)
4905{
4906 return PyModuleDef_Init(&itertoolsmodule);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004907}