blob: 293735a886428b5ab8ffdace03a4797089120ae9 [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;
866 PyObject *it = NULL;
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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000876 to = PyObject_GC_New(teeobject, &tee_type);
877 if (to == NULL)
878 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000879 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000880 if (!to->dataobj) {
881 PyObject_GC_Del(to);
882 to = NULL;
883 goto done;
884 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000885
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000886 to->index = 0;
887 to->weakreflist = NULL;
888 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000889done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000890 Py_XDECREF(it);
891 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000892}
893
Tal Einatc4bccd32018-09-12 00:49:13 +0300894/*[clinic input]
895@classmethod
896itertools._tee.__new__
897 iterable: object
898 /
899Iterator wrapped to make it copyable.
900[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000901
Tal Einatc4bccd32018-09-12 00:49:13 +0300902static PyObject *
903itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
904/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
905{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000906 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000907}
908
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000909static int
910tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000911{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000912 if (to->weakreflist != NULL)
913 PyObject_ClearWeakRefs((PyObject *) to);
914 Py_CLEAR(to->dataobj);
915 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000916}
917
918static void
919tee_dealloc(teeobject *to)
920{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000921 PyObject_GC_UnTrack(to);
922 tee_clear(to);
923 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000924}
925
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000926static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530927tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000928{
929 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
930}
931
932static PyObject *
933tee_setstate(teeobject *to, PyObject *state)
934{
935 teedataobject *tdo;
936 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300937 if (!PyTuple_Check(state)) {
938 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000939 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300940 }
941 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
942 return NULL;
943 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000944 if (index < 0 || index > LINKCELLS) {
945 PyErr_SetString(PyExc_ValueError, "Index out of range");
946 return NULL;
947 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200948 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300949 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000950 to->index = index;
951 Py_RETURN_NONE;
952}
953
Raymond Hettingerad983e72003-11-12 14:32:26 +0000954static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700955 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
956 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
957 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000958 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000959};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000960
961static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000962 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000963 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000964 sizeof(teeobject), /* tp_basicsize */
965 0, /* tp_itemsize */
966 /* methods */
967 (destructor)tee_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200968 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000969 0, /* tp_getattr */
970 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200971 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000972 0, /* tp_repr */
973 0, /* tp_as_number */
974 0, /* tp_as_sequence */
975 0, /* tp_as_mapping */
976 0, /* tp_hash */
977 0, /* tp_call */
978 0, /* tp_str */
979 0, /* tp_getattro */
980 0, /* tp_setattro */
981 0, /* tp_as_buffer */
982 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300983 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000984 (traverseproc)tee_traverse, /* tp_traverse */
985 (inquiry)tee_clear, /* tp_clear */
986 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700987 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000988 PyObject_SelfIter, /* tp_iter */
989 (iternextfunc)tee_next, /* tp_iternext */
990 tee_methods, /* tp_methods */
991 0, /* tp_members */
992 0, /* tp_getset */
993 0, /* tp_base */
994 0, /* tp_dict */
995 0, /* tp_descr_get */
996 0, /* tp_descr_set */
997 0, /* tp_dictoffset */
998 0, /* tp_init */
999 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001000 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001001 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001002};
1003
Tal Einatc4bccd32018-09-12 00:49:13 +03001004/*[clinic input]
1005itertools.tee
1006 iterable: object
1007 n: Py_ssize_t = 2
1008 /
1009Returns a tuple of n independent iterators.
1010[clinic start generated code]*/
1011
Raymond Hettingerad983e72003-11-12 14:32:26 +00001012static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001013itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
1014/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +00001015{
Tal Einatc4bccd32018-09-12 00:49:13 +03001016 Py_ssize_t i;
1017 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02001018 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +00001019
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001020 if (n < 0) {
1021 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
1022 return NULL;
1023 }
1024 result = PyTuple_New(n);
1025 if (result == NULL)
1026 return NULL;
1027 if (n == 0)
1028 return result;
1029 it = PyObject_GetIter(iterable);
1030 if (it == NULL) {
1031 Py_DECREF(result);
1032 return NULL;
1033 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001034
Serhiy Storchakaf320be72018-01-25 10:49:40 +02001035 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001036 Py_DECREF(it);
1037 Py_DECREF(result);
1038 return NULL;
1039 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +02001040 if (copyfunc != NULL) {
1041 copyable = it;
1042 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001043 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001044 copyable = tee_fromiterable(it);
1045 Py_DECREF(it);
1046 if (copyable == NULL) {
1047 Py_DECREF(result);
1048 return NULL;
1049 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001050 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
1051 if (copyfunc == NULL) {
1052 Py_DECREF(copyable);
1053 Py_DECREF(result);
1054 return NULL;
1055 }
1056 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02001057
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001058 PyTuple_SET_ITEM(result, 0, copyable);
1059 for (i = 1; i < n; i++) {
1060 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001061 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001062 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001063 Py_DECREF(result);
1064 return NULL;
1065 }
1066 PyTuple_SET_ITEM(result, i, copyable);
1067 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001068 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001069 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +00001070}
1071
Raymond Hettingerad983e72003-11-12 14:32:26 +00001072
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001073/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001074
1075typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001076 PyObject_HEAD
1077 PyObject *it;
1078 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001079 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001080 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001081} cycleobject;
1082
Tal Einatc4bccd32018-09-12 00:49:13 +03001083/*[clinic input]
1084@classmethod
1085itertools.cycle.__new__
1086 iterable: object
1087 /
1088Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
1089[clinic start generated code]*/
1090
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001091static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001092itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
1093/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001094{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001095 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001096 PyObject *saved;
1097 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001098
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001099 /* Get iterator. */
1100 it = PyObject_GetIter(iterable);
1101 if (it == NULL)
1102 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001103
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001104 saved = PyList_New(0);
1105 if (saved == NULL) {
1106 Py_DECREF(it);
1107 return NULL;
1108 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001109
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001110 /* create cycleobject structure */
1111 lz = (cycleobject *)type->tp_alloc(type, 0);
1112 if (lz == NULL) {
1113 Py_DECREF(it);
1114 Py_DECREF(saved);
1115 return NULL;
1116 }
1117 lz->it = it;
1118 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001119 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001122 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001123}
1124
1125static void
1126cycle_dealloc(cycleobject *lz)
1127{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001128 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001129 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001130 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001131 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001132}
1133
1134static int
1135cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1136{
Hai Shi47a23fc2020-06-07 20:05:36 +08001137 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001138 Py_VISIT(lz->saved);
1139 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001140}
1141
1142static PyObject *
1143cycle_next(cycleobject *lz)
1144{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001145 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001146
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001147 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001148 item = PyIter_Next(lz->it);
1149 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001150 if (lz->firstpass)
1151 return item;
1152 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001153 Py_DECREF(item);
1154 return NULL;
1155 }
1156 return item;
1157 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001158 /* Note: StopIteration is already cleared by PyIter_Next() */
1159 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001160 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001161 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001162 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001163 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001164 return NULL;
1165 item = PyList_GET_ITEM(lz->saved, lz->index);
1166 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001167 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001168 lz->index = 0;
1169 Py_INCREF(item);
1170 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001171}
1172
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001173static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301174cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001175{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001176 /* Create a new cycle with the iterator tuple, then set the saved state */
1177 if (lz->it == NULL) {
1178 PyObject *it = PyObject_GetIter(lz->saved);
1179 if (it == NULL)
1180 return NULL;
1181 if (lz->index != 0) {
1182 _Py_IDENTIFIER(__setstate__);
1183 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1184 "n", lz->index);
1185 if (res == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 Py_DECREF(res);
1190 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001191 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001192 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001193 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1194 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001195}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001196
1197static PyObject *
1198cycle_setstate(cycleobject *lz, PyObject *state)
1199{
1200 PyObject *saved=NULL;
1201 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001202 if (!PyTuple_Check(state)) {
1203 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001204 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001205 }
1206 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1207 return NULL;
1208 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001209 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001210 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001211 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001212 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001213 Py_RETURN_NONE;
1214}
1215
1216static PyMethodDef cycle_methods[] = {
1217 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1218 reduce_doc},
1219 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1220 setstate_doc},
1221 {NULL, NULL} /* sentinel */
1222};
1223
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001224static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001225 PyVarObject_HEAD_INIT(NULL, 0)
1226 "itertools.cycle", /* tp_name */
1227 sizeof(cycleobject), /* tp_basicsize */
1228 0, /* tp_itemsize */
1229 /* methods */
1230 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001231 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001232 0, /* tp_getattr */
1233 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001234 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001235 0, /* tp_repr */
1236 0, /* tp_as_number */
1237 0, /* tp_as_sequence */
1238 0, /* tp_as_mapping */
1239 0, /* tp_hash */
1240 0, /* tp_call */
1241 0, /* tp_str */
1242 PyObject_GenericGetAttr, /* tp_getattro */
1243 0, /* tp_setattro */
1244 0, /* tp_as_buffer */
1245 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1246 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001247 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001248 (traverseproc)cycle_traverse, /* tp_traverse */
1249 0, /* tp_clear */
1250 0, /* tp_richcompare */
1251 0, /* tp_weaklistoffset */
1252 PyObject_SelfIter, /* tp_iter */
1253 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001254 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001255 0, /* tp_members */
1256 0, /* tp_getset */
1257 0, /* tp_base */
1258 0, /* tp_dict */
1259 0, /* tp_descr_get */
1260 0, /* tp_descr_set */
1261 0, /* tp_dictoffset */
1262 0, /* tp_init */
1263 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001264 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001265 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001266};
1267
1268
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269/* dropwhile object **********************************************************/
1270
1271typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001272 PyObject_HEAD
1273 PyObject *func;
1274 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001275 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001276} dropwhileobject;
1277
Tal Einatc4bccd32018-09-12 00:49:13 +03001278/*[clinic input]
1279@classmethod
1280itertools.dropwhile.__new__
1281 predicate as func: object
1282 iterable as seq: object
1283 /
1284Drop items from the iterable while predicate(item) is true.
1285
1286Afterwards, return every element until the iterable is exhausted.
1287[clinic start generated code]*/
1288
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001289static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001290itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1291/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001293 PyObject *it;
1294 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001296 /* Get iterator. */
1297 it = PyObject_GetIter(seq);
1298 if (it == NULL)
1299 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001301 /* create dropwhileobject structure */
1302 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1303 if (lz == NULL) {
1304 Py_DECREF(it);
1305 return NULL;
1306 }
1307 Py_INCREF(func);
1308 lz->func = func;
1309 lz->it = it;
1310 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001312 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001313}
1314
1315static void
1316dropwhile_dealloc(dropwhileobject *lz)
1317{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001318 PyObject_GC_UnTrack(lz);
1319 Py_XDECREF(lz->func);
1320 Py_XDECREF(lz->it);
1321 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001322}
1323
1324static int
1325dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1326{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001327 Py_VISIT(lz->it);
1328 Py_VISIT(lz->func);
1329 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001330}
1331
1332static PyObject *
1333dropwhile_next(dropwhileobject *lz)
1334{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001335 PyObject *item, *good;
1336 PyObject *it = lz->it;
1337 long ok;
1338 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001340 iternext = *Py_TYPE(it)->tp_iternext;
1341 for (;;) {
1342 item = iternext(it);
1343 if (item == NULL)
1344 return NULL;
1345 if (lz->start == 1)
1346 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001347
Petr Viktorinffd97532020-02-11 17:46:57 +01001348 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001349 if (good == NULL) {
1350 Py_DECREF(item);
1351 return NULL;
1352 }
1353 ok = PyObject_IsTrue(good);
1354 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001355 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001356 lz->start = 1;
1357 return item;
1358 }
1359 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001360 if (ok < 0)
1361 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001362 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363}
1364
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001365static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301366dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001367{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001368 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 +00001369}
1370
1371static PyObject *
1372dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1373{
1374 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001375 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001376 return NULL;
1377 lz->start = start;
1378 Py_RETURN_NONE;
1379}
1380
1381static PyMethodDef dropwhile_methods[] = {
1382 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1383 reduce_doc},
1384 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1385 setstate_doc},
1386 {NULL, NULL} /* sentinel */
1387};
1388
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001389static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001390 PyVarObject_HEAD_INIT(NULL, 0)
1391 "itertools.dropwhile", /* tp_name */
1392 sizeof(dropwhileobject), /* tp_basicsize */
1393 0, /* tp_itemsize */
1394 /* methods */
1395 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001396 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001397 0, /* tp_getattr */
1398 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001399 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001400 0, /* tp_repr */
1401 0, /* tp_as_number */
1402 0, /* tp_as_sequence */
1403 0, /* tp_as_mapping */
1404 0, /* tp_hash */
1405 0, /* tp_call */
1406 0, /* tp_str */
1407 PyObject_GenericGetAttr, /* tp_getattro */
1408 0, /* tp_setattro */
1409 0, /* tp_as_buffer */
1410 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1411 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001412 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001413 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001414 0, /* tp_clear */
1415 0, /* tp_richcompare */
1416 0, /* tp_weaklistoffset */
1417 PyObject_SelfIter, /* tp_iter */
1418 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001419 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001420 0, /* tp_members */
1421 0, /* tp_getset */
1422 0, /* tp_base */
1423 0, /* tp_dict */
1424 0, /* tp_descr_get */
1425 0, /* tp_descr_set */
1426 0, /* tp_dictoffset */
1427 0, /* tp_init */
1428 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001429 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001430 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001431};
1432
1433
1434/* takewhile object **********************************************************/
1435
1436typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 PyObject_HEAD
1438 PyObject *func;
1439 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001440 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001441} takewhileobject;
1442
Tal Einatc4bccd32018-09-12 00:49:13 +03001443/*[clinic input]
1444@classmethod
1445itertools.takewhile.__new__
1446 predicate as func: object
1447 iterable as seq: object
1448 /
1449Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1450[clinic start generated code]*/
1451
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001453itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1454/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 PyObject *it;
1457 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001459 /* Get iterator. */
1460 it = PyObject_GetIter(seq);
1461 if (it == NULL)
1462 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001464 /* create takewhileobject structure */
1465 lz = (takewhileobject *)type->tp_alloc(type, 0);
1466 if (lz == NULL) {
1467 Py_DECREF(it);
1468 return NULL;
1469 }
1470 Py_INCREF(func);
1471 lz->func = func;
1472 lz->it = it;
1473 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001475 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476}
1477
1478static void
1479takewhile_dealloc(takewhileobject *lz)
1480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001481 PyObject_GC_UnTrack(lz);
1482 Py_XDECREF(lz->func);
1483 Py_XDECREF(lz->it);
1484 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485}
1486
1487static int
1488takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001490 Py_VISIT(lz->it);
1491 Py_VISIT(lz->func);
1492 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001493}
1494
1495static PyObject *
1496takewhile_next(takewhileobject *lz)
1497{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 PyObject *item, *good;
1499 PyObject *it = lz->it;
1500 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001502 if (lz->stop == 1)
1503 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001505 item = (*Py_TYPE(it)->tp_iternext)(it);
1506 if (item == NULL)
1507 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001508
Petr Viktorinffd97532020-02-11 17:46:57 +01001509 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001510 if (good == NULL) {
1511 Py_DECREF(item);
1512 return NULL;
1513 }
1514 ok = PyObject_IsTrue(good);
1515 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001516 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001517 return item;
1518 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001519 if (ok == 0)
1520 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001521 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001522}
1523
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001524static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301525takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001526{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001527 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 +00001528}
1529
1530static PyObject *
1531takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1532{
1533 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001534
Antoine Pitrou721738f2012-08-15 23:20:39 +02001535 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001536 return NULL;
1537 lz->stop = stop;
1538 Py_RETURN_NONE;
1539}
1540
1541static PyMethodDef takewhile_reduce_methods[] = {
1542 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1543 reduce_doc},
1544 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1545 setstate_doc},
1546 {NULL, NULL} /* sentinel */
1547};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001548
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001549static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001550 PyVarObject_HEAD_INIT(NULL, 0)
1551 "itertools.takewhile", /* tp_name */
1552 sizeof(takewhileobject), /* tp_basicsize */
1553 0, /* tp_itemsize */
1554 /* methods */
1555 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001556 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001557 0, /* tp_getattr */
1558 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001559 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001560 0, /* tp_repr */
1561 0, /* tp_as_number */
1562 0, /* tp_as_sequence */
1563 0, /* tp_as_mapping */
1564 0, /* tp_hash */
1565 0, /* tp_call */
1566 0, /* tp_str */
1567 PyObject_GenericGetAttr, /* tp_getattro */
1568 0, /* tp_setattro */
1569 0, /* tp_as_buffer */
1570 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1571 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001572 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001573 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001574 0, /* tp_clear */
1575 0, /* tp_richcompare */
1576 0, /* tp_weaklistoffset */
1577 PyObject_SelfIter, /* tp_iter */
1578 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001579 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001580 0, /* tp_members */
1581 0, /* tp_getset */
1582 0, /* tp_base */
1583 0, /* tp_dict */
1584 0, /* tp_descr_get */
1585 0, /* tp_descr_set */
1586 0, /* tp_dictoffset */
1587 0, /* tp_init */
1588 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001589 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001590 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001591};
1592
1593
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001594/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595
1596typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001597 PyObject_HEAD
1598 PyObject *it;
1599 Py_ssize_t next;
1600 Py_ssize_t stop;
1601 Py_ssize_t step;
1602 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001603} isliceobject;
1604
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001605static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606
1607static PyObject *
1608islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1609{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001610 PyObject *seq;
1611 Py_ssize_t start=0, stop=-1, step=1;
1612 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1613 Py_ssize_t numargs;
1614 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001616 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001617 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001618
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001619 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1620 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001622 numargs = PyTuple_Size(args);
1623 if (numargs == 2) {
1624 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001625 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001626 if (stop == -1) {
1627 if (PyErr_Occurred())
1628 PyErr_Clear();
1629 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001630 "Stop argument for islice() must be None or "
1631 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001632 return NULL;
1633 }
1634 }
1635 } else {
1636 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001637 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001638 if (start == -1 && PyErr_Occurred())
1639 PyErr_Clear();
1640 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001641 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001642 if (stop == -1) {
1643 if (PyErr_Occurred())
1644 PyErr_Clear();
1645 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001646 "Stop argument for islice() must be None or "
1647 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001648 return NULL;
1649 }
1650 }
1651 }
1652 if (start<0 || stop<-1) {
1653 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001654 "Indices for islice() must be None or "
1655 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001656 return NULL;
1657 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001659 if (a3 != NULL) {
1660 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001661 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001662 if (step == -1 && PyErr_Occurred())
1663 PyErr_Clear();
1664 }
1665 if (step<1) {
1666 PyErr_SetString(PyExc_ValueError,
1667 "Step for islice() must be a positive integer or None.");
1668 return NULL;
1669 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001671 /* Get iterator. */
1672 it = PyObject_GetIter(seq);
1673 if (it == NULL)
1674 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001676 /* create isliceobject structure */
1677 lz = (isliceobject *)type->tp_alloc(type, 0);
1678 if (lz == NULL) {
1679 Py_DECREF(it);
1680 return NULL;
1681 }
1682 lz->it = it;
1683 lz->next = start;
1684 lz->stop = stop;
1685 lz->step = step;
1686 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001688 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689}
1690
1691static void
1692islice_dealloc(isliceobject *lz)
1693{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001694 PyObject_GC_UnTrack(lz);
1695 Py_XDECREF(lz->it);
1696 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697}
1698
1699static int
1700islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1701{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001702 Py_VISIT(lz->it);
1703 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001704}
1705
1706static PyObject *
1707islice_next(isliceobject *lz)
1708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001709 PyObject *item;
1710 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001711 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001712 Py_ssize_t oldnext;
1713 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001714
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001715 if (it == NULL)
1716 return NULL;
1717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001718 iternext = *Py_TYPE(it)->tp_iternext;
1719 while (lz->cnt < lz->next) {
1720 item = iternext(it);
1721 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001722 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001723 Py_DECREF(item);
1724 lz->cnt++;
1725 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001726 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001727 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001728 item = iternext(it);
1729 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001730 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001731 lz->cnt++;
1732 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001733 /* The (size_t) cast below avoids the danger of undefined
1734 behaviour from signed integer overflow. */
1735 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001736 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1737 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001738 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001739
1740empty:
1741 Py_CLEAR(lz->it);
1742 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001743}
1744
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001745static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301746islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001747{
1748 /* When unpickled, generate a new object with the same bounds,
1749 * then 'setstate' with the next and count
1750 */
1751 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001752
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001753 if (lz->it == NULL) {
1754 PyObject *empty_list;
1755 PyObject *empty_it;
1756 empty_list = PyList_New(0);
1757 if (empty_list == NULL)
1758 return NULL;
1759 empty_it = PyObject_GetIter(empty_list);
1760 Py_DECREF(empty_list);
1761 if (empty_it == NULL)
1762 return NULL;
1763 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1764 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001765 if (lz->stop == -1) {
1766 stop = Py_None;
1767 Py_INCREF(stop);
1768 } else {
1769 stop = PyLong_FromSsize_t(lz->stop);
1770 if (stop == NULL)
1771 return NULL;
1772 }
1773 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1774 lz->it, lz->next, stop, lz->step,
1775 lz->cnt);
1776}
1777
1778static PyObject *
1779islice_setstate(isliceobject *lz, PyObject *state)
1780{
1781 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001782
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001783 if (cnt == -1 && PyErr_Occurred())
1784 return NULL;
1785 lz->cnt = cnt;
1786 Py_RETURN_NONE;
1787}
1788
1789static PyMethodDef islice_methods[] = {
1790 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1791 reduce_doc},
1792 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1793 setstate_doc},
1794 {NULL, NULL} /* sentinel */
1795};
1796
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001797PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001798"islice(iterable, stop) --> islice object\n\
1799islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001800\n\
1801Return an iterator whose next() method returns selected values from an\n\
1802iterable. If start is specified, will skip all preceding elements;\n\
1803otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001804specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805skipped between successive calls. Works like a slice() on a list\n\
1806but returns an iterator.");
1807
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001808static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 PyVarObject_HEAD_INIT(NULL, 0)
1810 "itertools.islice", /* tp_name */
1811 sizeof(isliceobject), /* tp_basicsize */
1812 0, /* tp_itemsize */
1813 /* methods */
1814 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001815 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001816 0, /* tp_getattr */
1817 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001818 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001819 0, /* tp_repr */
1820 0, /* tp_as_number */
1821 0, /* tp_as_sequence */
1822 0, /* tp_as_mapping */
1823 0, /* tp_hash */
1824 0, /* tp_call */
1825 0, /* tp_str */
1826 PyObject_GenericGetAttr, /* tp_getattro */
1827 0, /* tp_setattro */
1828 0, /* tp_as_buffer */
1829 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1830 Py_TPFLAGS_BASETYPE, /* tp_flags */
1831 islice_doc, /* tp_doc */
1832 (traverseproc)islice_traverse, /* tp_traverse */
1833 0, /* tp_clear */
1834 0, /* tp_richcompare */
1835 0, /* tp_weaklistoffset */
1836 PyObject_SelfIter, /* tp_iter */
1837 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001838 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001839 0, /* tp_members */
1840 0, /* tp_getset */
1841 0, /* tp_base */
1842 0, /* tp_dict */
1843 0, /* tp_descr_get */
1844 0, /* tp_descr_set */
1845 0, /* tp_dictoffset */
1846 0, /* tp_init */
1847 0, /* tp_alloc */
1848 islice_new, /* tp_new */
1849 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001850};
1851
1852
1853/* starmap object ************************************************************/
1854
1855typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001856 PyObject_HEAD
1857 PyObject *func;
1858 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001859} starmapobject;
1860
Tal Einatc4bccd32018-09-12 00:49:13 +03001861/*[clinic input]
1862@classmethod
1863itertools.starmap.__new__
1864 function as func: object
1865 iterable as seq: object
1866 /
1867Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1868[clinic start generated code]*/
1869
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001870static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001871itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1872/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001874 PyObject *it;
1875 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001877 /* Get iterator. */
1878 it = PyObject_GetIter(seq);
1879 if (it == NULL)
1880 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001881
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001882 /* create starmapobject structure */
1883 lz = (starmapobject *)type->tp_alloc(type, 0);
1884 if (lz == NULL) {
1885 Py_DECREF(it);
1886 return NULL;
1887 }
1888 Py_INCREF(func);
1889 lz->func = func;
1890 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001891
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001892 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001893}
1894
1895static void
1896starmap_dealloc(starmapobject *lz)
1897{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001898 PyObject_GC_UnTrack(lz);
1899 Py_XDECREF(lz->func);
1900 Py_XDECREF(lz->it);
1901 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001902}
1903
1904static int
1905starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1906{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001907 Py_VISIT(lz->it);
1908 Py_VISIT(lz->func);
1909 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001910}
1911
1912static PyObject *
1913starmap_next(starmapobject *lz)
1914{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001915 PyObject *args;
1916 PyObject *result;
1917 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001918
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001919 args = (*Py_TYPE(it)->tp_iternext)(it);
1920 if (args == NULL)
1921 return NULL;
1922 if (!PyTuple_CheckExact(args)) {
1923 PyObject *newargs = PySequence_Tuple(args);
1924 Py_DECREF(args);
1925 if (newargs == NULL)
1926 return NULL;
1927 args = newargs;
1928 }
1929 result = PyObject_Call(lz->func, args, NULL);
1930 Py_DECREF(args);
1931 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001932}
1933
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001934static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301935starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001936{
1937 /* Just pickle the iterator */
1938 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1939}
1940
1941static PyMethodDef starmap_methods[] = {
1942 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1943 reduce_doc},
1944 {NULL, NULL} /* sentinel */
1945};
1946
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001947static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001948 PyVarObject_HEAD_INIT(NULL, 0)
1949 "itertools.starmap", /* tp_name */
1950 sizeof(starmapobject), /* tp_basicsize */
1951 0, /* tp_itemsize */
1952 /* methods */
1953 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001954 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001955 0, /* tp_getattr */
1956 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001957 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001958 0, /* tp_repr */
1959 0, /* tp_as_number */
1960 0, /* tp_as_sequence */
1961 0, /* tp_as_mapping */
1962 0, /* tp_hash */
1963 0, /* tp_call */
1964 0, /* tp_str */
1965 PyObject_GenericGetAttr, /* tp_getattro */
1966 0, /* tp_setattro */
1967 0, /* tp_as_buffer */
1968 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1969 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001970 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001971 (traverseproc)starmap_traverse, /* tp_traverse */
1972 0, /* tp_clear */
1973 0, /* tp_richcompare */
1974 0, /* tp_weaklistoffset */
1975 PyObject_SelfIter, /* tp_iter */
1976 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001977 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001978 0, /* tp_members */
1979 0, /* tp_getset */
1980 0, /* tp_base */
1981 0, /* tp_dict */
1982 0, /* tp_descr_get */
1983 0, /* tp_descr_set */
1984 0, /* tp_dictoffset */
1985 0, /* tp_init */
1986 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001987 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001988 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001989};
1990
1991
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001992/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001993
1994typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001995 PyObject_HEAD
1996 PyObject *source; /* Iterator over input iterables */
1997 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001998} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001999
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002000static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002001
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002002static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00002003chain_new_internal(PyTypeObject *type, PyObject *source)
2004{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002005 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00002006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002007 lz = (chainobject *)type->tp_alloc(type, 0);
2008 if (lz == NULL) {
2009 Py_DECREF(source);
2010 return NULL;
2011 }
2012
2013 lz->source = source;
2014 lz->active = NULL;
2015 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00002016}
2017
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002018static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002019chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002020{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002021 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002022
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03002023 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002024 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002025
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002026 source = PyObject_GetIter(args);
2027 if (source == NULL)
2028 return NULL;
2029
2030 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00002031}
2032
Tal Einatc4bccd32018-09-12 00:49:13 +03002033/*[clinic input]
2034@classmethod
2035itertools.chain.from_iterable
2036 iterable as arg: object
2037 /
2038Alternative chain() constructor taking a single iterable argument that evaluates lazily.
2039[clinic start generated code]*/
2040
Christian Heimesf16baeb2008-02-29 14:57:44 +00002041static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002042itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
2043/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00002044{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002045 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002046
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002047 source = PyObject_GetIter(arg);
2048 if (source == NULL)
2049 return NULL;
2050
2051 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002052}
2053
2054static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002055chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002056{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002057 PyObject_GC_UnTrack(lz);
2058 Py_XDECREF(lz->active);
2059 Py_XDECREF(lz->source);
2060 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002061}
2062
Raymond Hettinger2012f172003-02-07 05:32:58 +00002063static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002064chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00002065{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002066 Py_VISIT(lz->source);
2067 Py_VISIT(lz->active);
2068 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002069}
2070
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002071static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002072chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002073{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002074 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002075
T. Wouters5466d4a2017-03-30 09:58:35 -07002076 /* lz->source is the iterator of iterables. If it's NULL, we've already
2077 * consumed them all. lz->active is the current iterator. If it's NULL,
2078 * we should grab a new one from lz->source. */
2079 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002080 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07002081 PyObject *iterable = PyIter_Next(lz->source);
2082 if (iterable == NULL) {
2083 Py_CLEAR(lz->source);
2084 return NULL; /* no more input sources */
2085 }
2086 lz->active = PyObject_GetIter(iterable);
2087 Py_DECREF(iterable);
2088 if (lz->active == NULL) {
2089 Py_CLEAR(lz->source);
2090 return NULL; /* input not iterable */
2091 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002092 }
T. Wouters5466d4a2017-03-30 09:58:35 -07002093 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
2094 if (item != NULL)
2095 return item;
2096 if (PyErr_Occurred()) {
2097 if (PyErr_ExceptionMatches(PyExc_StopIteration))
2098 PyErr_Clear();
2099 else
2100 return NULL; /* input raised an exception */
2101 }
2102 /* lz->active is consumed, try with the next iterable. */
2103 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002104 }
T. Wouters5466d4a2017-03-30 09:58:35 -07002105 /* Everything had been consumed already. */
2106 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002107}
2108
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002109static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302110chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002111{
2112 if (lz->source) {
2113 /* we can't pickle function objects (itertools.from_iterable) so
2114 * we must use setstate to replace the iterable. One day we
2115 * will fix pickling of functions
2116 */
2117 if (lz->active) {
2118 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
2119 } else {
2120 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
2121 }
2122 } else {
2123 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
2124 }
2125 return NULL;
2126}
2127
2128static PyObject *
2129chain_setstate(chainobject *lz, PyObject *state)
2130{
2131 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002132
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002133 if (!PyTuple_Check(state)) {
2134 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002135 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002136 }
2137 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2138 return NULL;
2139 }
2140 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2141 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2142 return NULL;
2143 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002144
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002145 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002146 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002147 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002148 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002149 Py_RETURN_NONE;
2150}
2151
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002152PyDoc_STRVAR(chain_doc,
2153"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002154\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002155Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002156first iterable until it is exhausted, then elements from the next\n\
2157iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002158
Christian Heimesf16baeb2008-02-29 14:57:44 +00002159static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002160 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002161 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2162 reduce_doc},
2163 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2164 setstate_doc},
Ethan Smitha8403d02020-04-09 20:28:08 -07002165 {"__class_getitem__", (PyCFunction)Py_GenericAlias,
2166 METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002167 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002168};
2169
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002170static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002171 PyVarObject_HEAD_INIT(NULL, 0)
2172 "itertools.chain", /* tp_name */
2173 sizeof(chainobject), /* tp_basicsize */
2174 0, /* tp_itemsize */
2175 /* methods */
2176 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002177 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002178 0, /* tp_getattr */
2179 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002180 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002181 0, /* tp_repr */
2182 0, /* tp_as_number */
2183 0, /* tp_as_sequence */
2184 0, /* tp_as_mapping */
2185 0, /* tp_hash */
2186 0, /* tp_call */
2187 0, /* tp_str */
2188 PyObject_GenericGetAttr, /* tp_getattro */
2189 0, /* tp_setattro */
2190 0, /* tp_as_buffer */
2191 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2192 Py_TPFLAGS_BASETYPE, /* tp_flags */
2193 chain_doc, /* tp_doc */
2194 (traverseproc)chain_traverse, /* tp_traverse */
2195 0, /* tp_clear */
2196 0, /* tp_richcompare */
2197 0, /* tp_weaklistoffset */
2198 PyObject_SelfIter, /* tp_iter */
2199 (iternextfunc)chain_next, /* tp_iternext */
2200 chain_methods, /* tp_methods */
2201 0, /* tp_members */
2202 0, /* tp_getset */
2203 0, /* tp_base */
2204 0, /* tp_dict */
2205 0, /* tp_descr_get */
2206 0, /* tp_descr_set */
2207 0, /* tp_dictoffset */
2208 0, /* tp_init */
2209 0, /* tp_alloc */
2210 chain_new, /* tp_new */
2211 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002212};
2213
2214
Christian Heimesc3f30c42008-02-22 16:37:40 +00002215/* product object ************************************************************/
2216
2217typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002218 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002219 PyObject *pools; /* tuple of pool tuples */
2220 Py_ssize_t *indices; /* one index per pool */
2221 PyObject *result; /* most recently returned result tuple */
2222 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002223} productobject;
2224
2225static PyTypeObject product_type;
2226
2227static PyObject *
2228product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2229{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002230 productobject *lz;
2231 Py_ssize_t nargs, npools, repeat=1;
2232 PyObject *pools = NULL;
2233 Py_ssize_t *indices = NULL;
2234 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002235
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002236 if (kwds != NULL) {
2237 char *kwlist[] = {"repeat", 0};
2238 PyObject *tmpargs = PyTuple_New(0);
2239 if (tmpargs == NULL)
2240 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002241 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2242 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002243 Py_DECREF(tmpargs);
2244 return NULL;
2245 }
2246 Py_DECREF(tmpargs);
2247 if (repeat < 0) {
2248 PyErr_SetString(PyExc_ValueError,
2249 "repeat argument cannot be negative");
2250 return NULL;
2251 }
2252 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002253
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002254 assert(PyTuple_CheckExact(args));
2255 if (repeat == 0) {
2256 nargs = 0;
2257 } else {
2258 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002259 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002260 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2261 return NULL;
2262 }
2263 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002264 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002265
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002266 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002267 if (indices == NULL) {
2268 PyErr_NoMemory();
2269 goto error;
2270 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002271
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002272 pools = PyTuple_New(npools);
2273 if (pools == NULL)
2274 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002275
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002276 for (i=0; i < nargs ; ++i) {
2277 PyObject *item = PyTuple_GET_ITEM(args, i);
2278 PyObject *pool = PySequence_Tuple(item);
2279 if (pool == NULL)
2280 goto error;
2281 PyTuple_SET_ITEM(pools, i, pool);
2282 indices[i] = 0;
2283 }
2284 for ( ; i < npools; ++i) {
2285 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2286 Py_INCREF(pool);
2287 PyTuple_SET_ITEM(pools, i, pool);
2288 indices[i] = 0;
2289 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002290
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002291 /* create productobject structure */
2292 lz = (productobject *)type->tp_alloc(type, 0);
2293 if (lz == NULL)
2294 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002295
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002296 lz->pools = pools;
2297 lz->indices = indices;
2298 lz->result = NULL;
2299 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002300
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002301 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002302
2303error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002304 if (indices != NULL)
2305 PyMem_Free(indices);
2306 Py_XDECREF(pools);
2307 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002308}
2309
2310static void
2311product_dealloc(productobject *lz)
2312{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002313 PyObject_GC_UnTrack(lz);
2314 Py_XDECREF(lz->pools);
2315 Py_XDECREF(lz->result);
2316 if (lz->indices != NULL)
2317 PyMem_Free(lz->indices);
2318 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002319}
2320
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002321static PyObject *
2322product_sizeof(productobject *lz, void *unused)
2323{
2324 Py_ssize_t res;
2325
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002326 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002327 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2328 return PyLong_FromSsize_t(res);
2329}
2330
2331PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2332
Christian Heimesc3f30c42008-02-22 16:37:40 +00002333static int
2334product_traverse(productobject *lz, visitproc visit, void *arg)
2335{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002336 Py_VISIT(lz->pools);
2337 Py_VISIT(lz->result);
2338 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002339}
2340
2341static PyObject *
2342product_next(productobject *lz)
2343{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002344 PyObject *pool;
2345 PyObject *elem;
2346 PyObject *oldelem;
2347 PyObject *pools = lz->pools;
2348 PyObject *result = lz->result;
2349 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2350 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002351
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002352 if (lz->stopped)
2353 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002354
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002355 if (result == NULL) {
2356 /* On the first pass, return an initial tuple filled with the
2357 first element from each pool. */
2358 result = PyTuple_New(npools);
2359 if (result == NULL)
2360 goto empty;
2361 lz->result = result;
2362 for (i=0; i < npools; i++) {
2363 pool = PyTuple_GET_ITEM(pools, i);
2364 if (PyTuple_GET_SIZE(pool) == 0)
2365 goto empty;
2366 elem = PyTuple_GET_ITEM(pool, 0);
2367 Py_INCREF(elem);
2368 PyTuple_SET_ITEM(result, i, elem);
2369 }
2370 } else {
2371 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002372
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002373 /* Copy the previous result tuple or re-use it if available */
2374 if (Py_REFCNT(result) > 1) {
2375 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002376 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002377 if (result == NULL)
2378 goto empty;
2379 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002380 Py_DECREF(old_result);
2381 }
Brandt Bucher226a0122020-12-04 19:45:57 -08002382 // bpo-42536: The GC may have untracked this result tuple. Since we're
2383 // recycling it, make sure it's tracked again:
2384 else if (!_PyObject_GC_IS_TRACKED(result)) {
2385 _PyObject_GC_TRACK(result);
2386 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002387 /* Now, we've got the only copy so we can update it in-place */
2388 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002389
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002390 /* Update the pool indices right-to-left. Only advance to the
2391 next pool when the previous one rolls-over */
2392 for (i=npools-1 ; i >= 0 ; i--) {
2393 pool = PyTuple_GET_ITEM(pools, i);
2394 indices[i]++;
2395 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2396 /* Roll-over and advance to next pool */
2397 indices[i] = 0;
2398 elem = PyTuple_GET_ITEM(pool, 0);
2399 Py_INCREF(elem);
2400 oldelem = PyTuple_GET_ITEM(result, i);
2401 PyTuple_SET_ITEM(result, i, elem);
2402 Py_DECREF(oldelem);
2403 } else {
2404 /* No rollover. Just increment and stop here. */
2405 elem = PyTuple_GET_ITEM(pool, indices[i]);
2406 Py_INCREF(elem);
2407 oldelem = PyTuple_GET_ITEM(result, i);
2408 PyTuple_SET_ITEM(result, i, elem);
2409 Py_DECREF(oldelem);
2410 break;
2411 }
2412 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002413
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002414 /* If i is negative, then the indices have all rolled-over
2415 and we're done. */
2416 if (i < 0)
2417 goto empty;
2418 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002419
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002420 Py_INCREF(result);
2421 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002422
2423empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002424 lz->stopped = 1;
2425 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002426}
2427
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002428static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302429product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002430{
2431 if (lz->stopped) {
2432 return Py_BuildValue("O(())", Py_TYPE(lz));
2433 } else if (lz->result == NULL) {
2434 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2435 } else {
2436 PyObject *indices;
2437 Py_ssize_t n, i;
2438
2439 /* we must pickle the indices use them for setstate, and
2440 * additionally indicate that the iterator has started
2441 */
2442 n = PyTuple_GET_SIZE(lz->pools);
2443 indices = PyTuple_New(n);
2444 if (indices == NULL)
2445 return NULL;
2446 for (i=0; i<n; i++){
2447 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2448 if (!index) {
2449 Py_DECREF(indices);
2450 return NULL;
2451 }
2452 PyTuple_SET_ITEM(indices, i, index);
2453 }
2454 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2455 }
2456}
2457
2458static PyObject *
2459product_setstate(productobject *lz, PyObject *state)
2460{
2461 PyObject *result;
2462 Py_ssize_t n, i;
2463
2464 n = PyTuple_GET_SIZE(lz->pools);
2465 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2466 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2467 return NULL;
2468 }
2469 for (i=0; i<n; i++)
2470 {
2471 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2472 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002473 PyObject* pool;
2474 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002475 if (index < 0 && PyErr_Occurred())
2476 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002477 pool = PyTuple_GET_ITEM(lz->pools, i);
2478 poolsize = PyTuple_GET_SIZE(pool);
2479 if (poolsize == 0) {
2480 lz->stopped = 1;
2481 Py_RETURN_NONE;
2482 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002483 /* clamp the index */
2484 if (index < 0)
2485 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002486 else if (index > poolsize-1)
2487 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002488 lz->indices[i] = index;
2489 }
2490
2491 result = PyTuple_New(n);
2492 if (!result)
2493 return NULL;
2494 for (i=0; i<n; i++) {
2495 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2496 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2497 Py_INCREF(element);
2498 PyTuple_SET_ITEM(result, i, element);
2499 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002500 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002501 Py_RETURN_NONE;
2502}
2503
2504static PyMethodDef product_methods[] = {
2505 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2506 reduce_doc},
2507 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2508 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002509 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2510 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002511 {NULL, NULL} /* sentinel */
2512};
2513
Christian Heimesc3f30c42008-02-22 16:37:40 +00002514PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002515"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002516\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002517Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002518For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2519The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2520cycle in a manner similar to an odometer (with the rightmost element changing\n\
2521on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002522To compute the product of an iterable with itself, specify the number\n\
2523of repetitions with the optional repeat keyword argument. For example,\n\
2524product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002525product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2526product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2527
2528static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002529 PyVarObject_HEAD_INIT(NULL, 0)
2530 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002531 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002532 0, /* tp_itemsize */
2533 /* methods */
2534 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002535 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002536 0, /* tp_getattr */
2537 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002538 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002539 0, /* tp_repr */
2540 0, /* tp_as_number */
2541 0, /* tp_as_sequence */
2542 0, /* tp_as_mapping */
2543 0, /* tp_hash */
2544 0, /* tp_call */
2545 0, /* tp_str */
2546 PyObject_GenericGetAttr, /* tp_getattro */
2547 0, /* tp_setattro */
2548 0, /* tp_as_buffer */
2549 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2550 Py_TPFLAGS_BASETYPE, /* tp_flags */
2551 product_doc, /* tp_doc */
2552 (traverseproc)product_traverse, /* tp_traverse */
2553 0, /* tp_clear */
2554 0, /* tp_richcompare */
2555 0, /* tp_weaklistoffset */
2556 PyObject_SelfIter, /* tp_iter */
2557 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002558 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002559 0, /* tp_members */
2560 0, /* tp_getset */
2561 0, /* tp_base */
2562 0, /* tp_dict */
2563 0, /* tp_descr_get */
2564 0, /* tp_descr_set */
2565 0, /* tp_dictoffset */
2566 0, /* tp_init */
2567 0, /* tp_alloc */
2568 product_new, /* tp_new */
2569 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002570};
2571
2572
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002573/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002574
2575typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002576 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002577 PyObject *pool; /* input converted to a tuple */
2578 Py_ssize_t *indices; /* one index per result element */
2579 PyObject *result; /* most recently returned result tuple */
2580 Py_ssize_t r; /* size of result tuple */
2581 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002582} combinationsobject;
2583
Tal Einatc4bccd32018-09-12 00:49:13 +03002584
2585/*[clinic input]
2586@classmethod
2587itertools.combinations.__new__
2588 iterable: object
2589 r: Py_ssize_t
2590Return successive r-length combinations of elements in the iterable.
2591
2592combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2593[clinic start generated code]*/
2594
Christian Heimes380f7f22008-02-28 11:19:05 +00002595static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002596itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2597 Py_ssize_t r)
2598/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002599{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002600 combinationsobject *co;
2601 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002602 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002603 Py_ssize_t *indices = NULL;
2604 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002605
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002606 pool = PySequence_Tuple(iterable);
2607 if (pool == NULL)
2608 goto error;
2609 n = PyTuple_GET_SIZE(pool);
2610 if (r < 0) {
2611 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2612 goto error;
2613 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002614
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002615 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002616 if (indices == NULL) {
2617 PyErr_NoMemory();
2618 goto error;
2619 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002620
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002621 for (i=0 ; i<r ; i++)
2622 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002623
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002624 /* create combinationsobject structure */
2625 co = (combinationsobject *)type->tp_alloc(type, 0);
2626 if (co == NULL)
2627 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002629 co->pool = pool;
2630 co->indices = indices;
2631 co->result = NULL;
2632 co->r = r;
2633 co->stopped = r > n ? 1 : 0;
2634
2635 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002636
2637error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002638 if (indices != NULL)
2639 PyMem_Free(indices);
2640 Py_XDECREF(pool);
2641 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002642}
2643
2644static void
2645combinations_dealloc(combinationsobject *co)
2646{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002647 PyObject_GC_UnTrack(co);
2648 Py_XDECREF(co->pool);
2649 Py_XDECREF(co->result);
2650 if (co->indices != NULL)
2651 PyMem_Free(co->indices);
2652 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002653}
2654
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002655static PyObject *
2656combinations_sizeof(combinationsobject *co, void *unused)
2657{
2658 Py_ssize_t res;
2659
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002660 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002661 res += co->r * sizeof(Py_ssize_t);
2662 return PyLong_FromSsize_t(res);
2663}
2664
Christian Heimes380f7f22008-02-28 11:19:05 +00002665static int
2666combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2667{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002668 Py_VISIT(co->pool);
2669 Py_VISIT(co->result);
2670 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002671}
2672
2673static PyObject *
2674combinations_next(combinationsobject *co)
2675{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002676 PyObject *elem;
2677 PyObject *oldelem;
2678 PyObject *pool = co->pool;
2679 Py_ssize_t *indices = co->indices;
2680 PyObject *result = co->result;
2681 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2682 Py_ssize_t r = co->r;
2683 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002684
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002685 if (co->stopped)
2686 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002688 if (result == NULL) {
2689 /* On the first pass, initialize result tuple using the indices */
2690 result = PyTuple_New(r);
2691 if (result == NULL)
2692 goto empty;
2693 co->result = result;
2694 for (i=0; i<r ; i++) {
2695 index = indices[i];
2696 elem = PyTuple_GET_ITEM(pool, index);
2697 Py_INCREF(elem);
2698 PyTuple_SET_ITEM(result, i, elem);
2699 }
2700 } else {
2701 /* Copy the previous result tuple or re-use it if available */
2702 if (Py_REFCNT(result) > 1) {
2703 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002704 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002705 if (result == NULL)
2706 goto empty;
2707 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 Py_DECREF(old_result);
2709 }
Brandt Bucher226a0122020-12-04 19:45:57 -08002710 // bpo-42536: The GC may have untracked this result tuple. Since we're
2711 // recycling it, make sure it's tracked again:
2712 else if (!_PyObject_GC_IS_TRACKED(result)) {
2713 _PyObject_GC_TRACK(result);
2714 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002715 /* Now, we've got the only copy so we can update it in-place
2716 * CPython's empty tuple is a singleton and cached in
2717 * PyTuple's freelist.
2718 */
2719 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002720
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002721 /* Scan indices right-to-left until finding one that is not
2722 at its maximum (i + n - r). */
2723 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2724 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002725
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002726 /* If i is negative, then the indices are all at
2727 their maximum value and we're done. */
2728 if (i < 0)
2729 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002730
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002731 /* Increment the current index which we know is not at its
2732 maximum. Then move back to the right setting each index
2733 to its lowest possible value (one higher than the index
2734 to its left -- this maintains the sort order invariant). */
2735 indices[i]++;
2736 for (j=i+1 ; j<r ; j++)
2737 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002738
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002739 /* Update the result tuple for the new indices
2740 starting with i, the leftmost index that changed */
2741 for ( ; i<r ; i++) {
2742 index = indices[i];
2743 elem = PyTuple_GET_ITEM(pool, index);
2744 Py_INCREF(elem);
2745 oldelem = PyTuple_GET_ITEM(result, i);
2746 PyTuple_SET_ITEM(result, i, elem);
2747 Py_DECREF(oldelem);
2748 }
2749 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002751 Py_INCREF(result);
2752 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002753
2754empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002755 co->stopped = 1;
2756 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002757}
2758
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002759static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302760combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002761{
2762 if (lz->result == NULL) {
2763 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2764 } else if (lz->stopped) {
2765 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2766 } else {
2767 PyObject *indices;
2768 Py_ssize_t i;
2769
2770 /* we must pickle the indices and use them for setstate */
2771 indices = PyTuple_New(lz->r);
2772 if (!indices)
2773 return NULL;
2774 for (i=0; i<lz->r; i++)
2775 {
2776 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2777 if (!index) {
2778 Py_DECREF(indices);
2779 return NULL;
2780 }
2781 PyTuple_SET_ITEM(indices, i, index);
2782 }
2783
2784 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2785 }
2786}
2787
2788static PyObject *
2789combinations_setstate(combinationsobject *lz, PyObject *state)
2790{
2791 PyObject *result;
2792 Py_ssize_t i;
2793 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2794
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002795 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002796 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2797 return NULL;
2798 }
2799
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002800 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002801 Py_ssize_t max;
2802 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2803 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002804
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002805 if (index == -1 && PyErr_Occurred())
2806 return NULL; /* not an integer */
2807 max = i + n - lz->r;
2808 /* clamp the index (beware of negative max) */
2809 if (index > max)
2810 index = max;
2811 if (index < 0)
2812 index = 0;
2813 lz->indices[i] = index;
2814 }
2815
2816 result = PyTuple_New(lz->r);
2817 if (result == NULL)
2818 return NULL;
2819 for (i=0; i<lz->r; i++) {
2820 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2821 Py_INCREF(element);
2822 PyTuple_SET_ITEM(result, i, element);
2823 }
2824
Serhiy Storchaka48842712016-04-06 09:45:48 +03002825 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002826 Py_RETURN_NONE;
2827}
2828
2829static PyMethodDef combinations_methods[] = {
2830 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2831 reduce_doc},
2832 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2833 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002834 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2835 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002836 {NULL, NULL} /* sentinel */
2837};
2838
Christian Heimes380f7f22008-02-28 11:19:05 +00002839static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002840 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002841 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002842 sizeof(combinationsobject), /* tp_basicsize */
2843 0, /* tp_itemsize */
2844 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002845 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002846 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002847 0, /* tp_getattr */
2848 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002849 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002850 0, /* tp_repr */
2851 0, /* tp_as_number */
2852 0, /* tp_as_sequence */
2853 0, /* tp_as_mapping */
2854 0, /* tp_hash */
2855 0, /* tp_call */
2856 0, /* tp_str */
2857 PyObject_GenericGetAttr, /* tp_getattro */
2858 0, /* tp_setattro */
2859 0, /* tp_as_buffer */
2860 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2861 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002862 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002863 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002864 0, /* tp_clear */
2865 0, /* tp_richcompare */
2866 0, /* tp_weaklistoffset */
2867 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002868 (iternextfunc)combinations_next, /* tp_iternext */
2869 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002870 0, /* tp_members */
2871 0, /* tp_getset */
2872 0, /* tp_base */
2873 0, /* tp_dict */
2874 0, /* tp_descr_get */
2875 0, /* tp_descr_set */
2876 0, /* tp_dictoffset */
2877 0, /* tp_init */
2878 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002879 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002880 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002881};
2882
2883
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002884/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002885
2886/* Equivalent to:
2887
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002888 def combinations_with_replacement(iterable, r):
2889 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2890 # number items returned: (n+r-1)! / r! / (n-1)!
2891 pool = tuple(iterable)
2892 n = len(pool)
2893 indices = [0] * r
2894 yield tuple(pool[i] for i in indices)
2895 while 1:
2896 for i in reversed(range(r)):
2897 if indices[i] != n - 1:
2898 break
2899 else:
2900 return
2901 indices[i:] = [indices[i] + 1] * (r - i)
2902 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002903
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002904 def combinations_with_replacement2(iterable, r):
2905 'Alternate version that filters from product()'
2906 pool = tuple(iterable)
2907 n = len(pool)
2908 for indices in product(range(n), repeat=r):
2909 if sorted(indices) == list(indices):
2910 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002911*/
2912typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002913 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002914 PyObject *pool; /* input converted to a tuple */
2915 Py_ssize_t *indices; /* one index per result element */
2916 PyObject *result; /* most recently returned result tuple */
2917 Py_ssize_t r; /* size of result tuple */
2918 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002919} cwrobject;
2920
Tal Einatc4bccd32018-09-12 00:49:13 +03002921/*[clinic input]
2922@classmethod
2923itertools.combinations_with_replacement.__new__
2924 iterable: object
2925 r: Py_ssize_t
2926Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2927
2928combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2929[clinic start generated code]*/
2930
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002931static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002932itertools_combinations_with_replacement_impl(PyTypeObject *type,
2933 PyObject *iterable,
2934 Py_ssize_t r)
2935/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002936{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002937 cwrobject *co;
2938 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002939 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002940 Py_ssize_t *indices = NULL;
2941 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002942
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002943 pool = PySequence_Tuple(iterable);
2944 if (pool == NULL)
2945 goto error;
2946 n = PyTuple_GET_SIZE(pool);
2947 if (r < 0) {
2948 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2949 goto error;
2950 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002951
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002952 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002953 if (indices == NULL) {
2954 PyErr_NoMemory();
2955 goto error;
2956 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002957
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002958 for (i=0 ; i<r ; i++)
2959 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002960
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002961 /* create cwrobject structure */
2962 co = (cwrobject *)type->tp_alloc(type, 0);
2963 if (co == NULL)
2964 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002965
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002966 co->pool = pool;
2967 co->indices = indices;
2968 co->result = NULL;
2969 co->r = r;
2970 co->stopped = !n && r;
2971
2972 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002973
2974error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002975 if (indices != NULL)
2976 PyMem_Free(indices);
2977 Py_XDECREF(pool);
2978 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002979}
2980
2981static void
2982cwr_dealloc(cwrobject *co)
2983{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002984 PyObject_GC_UnTrack(co);
2985 Py_XDECREF(co->pool);
2986 Py_XDECREF(co->result);
2987 if (co->indices != NULL)
2988 PyMem_Free(co->indices);
2989 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002990}
2991
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002992static PyObject *
2993cwr_sizeof(cwrobject *co, void *unused)
2994{
2995 Py_ssize_t res;
2996
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002997 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002998 res += co->r * sizeof(Py_ssize_t);
2999 return PyLong_FromSsize_t(res);
3000}
3001
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003002static int
3003cwr_traverse(cwrobject *co, visitproc visit, void *arg)
3004{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003005 Py_VISIT(co->pool);
3006 Py_VISIT(co->result);
3007 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003008}
3009
3010static PyObject *
3011cwr_next(cwrobject *co)
3012{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003013 PyObject *elem;
3014 PyObject *oldelem;
3015 PyObject *pool = co->pool;
3016 Py_ssize_t *indices = co->indices;
3017 PyObject *result = co->result;
3018 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3019 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05003020 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003021
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003022 if (co->stopped)
3023 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003024
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05003026 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003027 result = PyTuple_New(r);
3028 if (result == NULL)
3029 goto empty;
3030 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07003031 if (n > 0) {
3032 elem = PyTuple_GET_ITEM(pool, 0);
3033 for (i=0; i<r ; i++) {
3034 assert(indices[i] == 0);
3035 Py_INCREF(elem);
3036 PyTuple_SET_ITEM(result, i, elem);
3037 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003038 }
3039 } else {
3040 /* Copy the previous result tuple or re-use it if available */
3041 if (Py_REFCNT(result) > 1) {
3042 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003043 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003044 if (result == NULL)
3045 goto empty;
3046 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003047 Py_DECREF(old_result);
3048 }
Brandt Bucher226a0122020-12-04 19:45:57 -08003049 // bpo-42536: The GC may have untracked this result tuple. Since we're
3050 // recycling it, make sure it's tracked again:
3051 else if (!_PyObject_GC_IS_TRACKED(result)) {
3052 _PyObject_GC_TRACK(result);
3053 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003054 /* Now, we've got the only copy so we can update it in-place CPython's
3055 empty tuple is a singleton and cached in PyTuple's freelist. */
3056 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003057
Tim Peters9edb1682013-09-03 11:49:31 -05003058 /* Scan indices right-to-left until finding one that is not
3059 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003060 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
3061 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003062
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003063 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05003064 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003065 if (i < 0)
3066 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003068 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05003069 maximum. Then set all to the right to the same value. */
3070 index = indices[i] + 1;
3071 assert(index < n);
3072 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003073 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05003074 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003075 Py_INCREF(elem);
3076 oldelem = PyTuple_GET_ITEM(result, i);
3077 PyTuple_SET_ITEM(result, i, elem);
3078 Py_DECREF(oldelem);
3079 }
3080 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003081
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003082 Py_INCREF(result);
3083 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003084
3085empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003086 co->stopped = 1;
3087 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003088}
3089
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003090static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303091cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003092{
3093 if (lz->result == NULL) {
3094 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
3095 } else if (lz->stopped) {
3096 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
3097 } else {
3098 PyObject *indices;
3099 Py_ssize_t i;
3100
3101 /* we must pickle the indices and use them for setstate */
3102 indices = PyTuple_New(lz->r);
3103 if (!indices)
3104 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003105 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003106 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
3107 if (!index) {
3108 Py_DECREF(indices);
3109 return NULL;
3110 }
3111 PyTuple_SET_ITEM(indices, i, index);
3112 }
3113
3114 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
3115 }
3116}
3117
3118static PyObject *
3119cwr_setstate(cwrobject *lz, PyObject *state)
3120{
3121 PyObject *result;
3122 Py_ssize_t n, i;
3123
3124 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
3125 {
3126 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3127 return NULL;
3128 }
3129
3130 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003131 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003132 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
3133 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003134
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003135 if (index < 0 && PyErr_Occurred())
3136 return NULL; /* not an integer */
3137 /* clamp the index */
3138 if (index < 0)
3139 index = 0;
3140 else if (index > n-1)
3141 index = n-1;
3142 lz->indices[i] = index;
3143 }
3144 result = PyTuple_New(lz->r);
3145 if (result == NULL)
3146 return NULL;
3147 for (i=0; i<lz->r; i++) {
3148 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3149 Py_INCREF(element);
3150 PyTuple_SET_ITEM(result, i, element);
3151 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003152 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003153 Py_RETURN_NONE;
3154}
3155
3156static PyMethodDef cwr_methods[] = {
3157 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3158 reduce_doc},
3159 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3160 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003161 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3162 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003163 {NULL, NULL} /* sentinel */
3164};
3165
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003166static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003167 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003168 "itertools.combinations_with_replacement", /* tp_name */
3169 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003170 0, /* tp_itemsize */
3171 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003172 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003173 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003174 0, /* tp_getattr */
3175 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003176 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003177 0, /* tp_repr */
3178 0, /* tp_as_number */
3179 0, /* tp_as_sequence */
3180 0, /* tp_as_mapping */
3181 0, /* tp_hash */
3182 0, /* tp_call */
3183 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003184 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003185 0, /* tp_setattro */
3186 0, /* tp_as_buffer */
3187 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003188 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003189 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003190 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003191 0, /* tp_clear */
3192 0, /* tp_richcompare */
3193 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003194 PyObject_SelfIter, /* tp_iter */
3195 (iternextfunc)cwr_next, /* tp_iternext */
3196 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003197 0, /* tp_members */
3198 0, /* tp_getset */
3199 0, /* tp_base */
3200 0, /* tp_dict */
3201 0, /* tp_descr_get */
3202 0, /* tp_descr_set */
3203 0, /* tp_dictoffset */
3204 0, /* tp_init */
3205 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003206 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003207 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003208};
3209
3210
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003211/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003212
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003213def permutations(iterable, r=None):
Sergey0078a0c2019-10-29 08:10:24 +03003214 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3215 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003216 pool = tuple(iterable)
3217 n = len(pool)
3218 r = n if r is None else r
Sergey0078a0c2019-10-29 08:10:24 +03003219 if r > n:
3220 return
3221 indices = list(range(n))
3222 cycles = list(range(n, n-r, -1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003223 yield tuple(pool[i] for i in indices[:r])
3224 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003225 for i in reversed(range(r)):
3226 cycles[i] -= 1
3227 if cycles[i] == 0:
3228 indices[i:] = indices[i+1:] + indices[i:i+1]
3229 cycles[i] = n - i
3230 else:
3231 j = cycles[i]
3232 indices[i], indices[-j] = indices[-j], indices[i]
3233 yield tuple(pool[i] for i in indices[:r])
3234 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003235 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003236 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003237*/
3238
3239typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003240 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003241 PyObject *pool; /* input converted to a tuple */
3242 Py_ssize_t *indices; /* one index per element in the pool */
3243 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3244 PyObject *result; /* most recently returned result tuple */
3245 Py_ssize_t r; /* size of result tuple */
3246 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003247} permutationsobject;
3248
Tal Einatc4bccd32018-09-12 00:49:13 +03003249/*[clinic input]
3250@classmethod
3251itertools.permutations.__new__
3252 iterable: object
3253 r as robj: object = None
3254Return successive r-length permutations of elements in the iterable.
3255
3256permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3257[clinic start generated code]*/
3258
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003259static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003260itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3261 PyObject *robj)
3262/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003263{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003264 permutationsobject *po;
3265 Py_ssize_t n;
3266 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003267 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003268 Py_ssize_t *indices = NULL;
3269 Py_ssize_t *cycles = NULL;
3270 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003271
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003272 pool = PySequence_Tuple(iterable);
3273 if (pool == NULL)
3274 goto error;
3275 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003276
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003277 r = n;
3278 if (robj != Py_None) {
3279 if (!PyLong_Check(robj)) {
3280 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3281 goto error;
3282 }
3283 r = PyLong_AsSsize_t(robj);
3284 if (r == -1 && PyErr_Occurred())
3285 goto error;
3286 }
3287 if (r < 0) {
3288 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3289 goto error;
3290 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003291
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003292 indices = PyMem_New(Py_ssize_t, n);
3293 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003294 if (indices == NULL || cycles == NULL) {
3295 PyErr_NoMemory();
3296 goto error;
3297 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003298
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003299 for (i=0 ; i<n ; i++)
3300 indices[i] = i;
3301 for (i=0 ; i<r ; i++)
3302 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003303
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003304 /* create permutationsobject structure */
3305 po = (permutationsobject *)type->tp_alloc(type, 0);
3306 if (po == NULL)
3307 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003308
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003309 po->pool = pool;
3310 po->indices = indices;
3311 po->cycles = cycles;
3312 po->result = NULL;
3313 po->r = r;
3314 po->stopped = r > n ? 1 : 0;
3315
3316 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003317
3318error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003319 if (indices != NULL)
3320 PyMem_Free(indices);
3321 if (cycles != NULL)
3322 PyMem_Free(cycles);
3323 Py_XDECREF(pool);
3324 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003325}
3326
3327static void
3328permutations_dealloc(permutationsobject *po)
3329{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003330 PyObject_GC_UnTrack(po);
3331 Py_XDECREF(po->pool);
3332 Py_XDECREF(po->result);
3333 PyMem_Free(po->indices);
3334 PyMem_Free(po->cycles);
3335 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003336}
3337
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003338static PyObject *
3339permutations_sizeof(permutationsobject *po, void *unused)
3340{
3341 Py_ssize_t res;
3342
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003343 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003344 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3345 res += po->r * sizeof(Py_ssize_t);
3346 return PyLong_FromSsize_t(res);
3347}
3348
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003349static int
3350permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3351{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003352 Py_VISIT(po->pool);
3353 Py_VISIT(po->result);
3354 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003355}
3356
3357static PyObject *
3358permutations_next(permutationsobject *po)
3359{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003360 PyObject *elem;
3361 PyObject *oldelem;
3362 PyObject *pool = po->pool;
3363 Py_ssize_t *indices = po->indices;
3364 Py_ssize_t *cycles = po->cycles;
3365 PyObject *result = po->result;
3366 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3367 Py_ssize_t r = po->r;
3368 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003369
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003370 if (po->stopped)
3371 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003372
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003373 if (result == NULL) {
3374 /* On the first pass, initialize result tuple using the indices */
3375 result = PyTuple_New(r);
3376 if (result == NULL)
3377 goto empty;
3378 po->result = result;
3379 for (i=0; i<r ; i++) {
3380 index = indices[i];
3381 elem = PyTuple_GET_ITEM(pool, index);
3382 Py_INCREF(elem);
3383 PyTuple_SET_ITEM(result, i, elem);
3384 }
3385 } else {
3386 if (n == 0)
3387 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003388
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003389 /* Copy the previous result tuple or re-use it if available */
3390 if (Py_REFCNT(result) > 1) {
3391 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003392 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003393 if (result == NULL)
3394 goto empty;
3395 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003396 Py_DECREF(old_result);
3397 }
Brandt Bucher226a0122020-12-04 19:45:57 -08003398 // bpo-42536: The GC may have untracked this result tuple. Since we're
3399 // recycling it, make sure it's tracked again:
3400 else if (!_PyObject_GC_IS_TRACKED(result)) {
3401 _PyObject_GC_TRACK(result);
3402 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003403 /* Now, we've got the only copy so we can update it in-place */
3404 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003405
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003406 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3407 for (i=r-1 ; i>=0 ; i--) {
3408 cycles[i] -= 1;
3409 if (cycles[i] == 0) {
3410 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3411 index = indices[i];
3412 for (j=i ; j<n-1 ; j++)
3413 indices[j] = indices[j+1];
3414 indices[n-1] = index;
3415 cycles[i] = n - i;
3416 } else {
3417 j = cycles[i];
3418 index = indices[i];
3419 indices[i] = indices[n-j];
3420 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003421
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003422 for (k=i; k<r ; k++) {
3423 /* start with i, the leftmost element that changed */
3424 /* yield tuple(pool[k] for k in indices[:r]) */
3425 index = indices[k];
3426 elem = PyTuple_GET_ITEM(pool, index);
3427 Py_INCREF(elem);
3428 oldelem = PyTuple_GET_ITEM(result, k);
3429 PyTuple_SET_ITEM(result, k, elem);
3430 Py_DECREF(oldelem);
3431 }
3432 break;
3433 }
3434 }
3435 /* If i is negative, then the cycles have all
3436 rolled-over and we're done. */
3437 if (i < 0)
3438 goto empty;
3439 }
3440 Py_INCREF(result);
3441 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003442
3443empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003444 po->stopped = 1;
3445 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003446}
3447
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003448static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303449permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003450{
3451 if (po->result == NULL) {
3452 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3453 } else if (po->stopped) {
3454 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3455 } else {
3456 PyObject *indices=NULL, *cycles=NULL;
3457 Py_ssize_t n, i;
3458
3459 /* we must pickle the indices and cycles and use them for setstate */
3460 n = PyTuple_GET_SIZE(po->pool);
3461 indices = PyTuple_New(n);
3462 if (indices == NULL)
3463 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003464 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003465 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3466 if (!index)
3467 goto err;
3468 PyTuple_SET_ITEM(indices, i, index);
3469 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003470
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003471 cycles = PyTuple_New(po->r);
3472 if (cycles == NULL)
3473 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003474 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003475 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3476 if (!index)
3477 goto err;
3478 PyTuple_SET_ITEM(cycles, i, index);
3479 }
3480 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3481 po->pool, po->r,
3482 indices, cycles);
3483 err:
3484 Py_XDECREF(indices);
3485 Py_XDECREF(cycles);
3486 return NULL;
3487 }
3488}
3489
3490static PyObject *
3491permutations_setstate(permutationsobject *po, PyObject *state)
3492{
3493 PyObject *indices, *cycles, *result;
3494 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003495
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003496 if (!PyTuple_Check(state)) {
3497 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3498 return NULL;
3499 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003500 if (!PyArg_ParseTuple(state, "O!O!",
3501 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003502 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003503 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003504 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003505
3506 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003507 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003508 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3509 return NULL;
3510 }
3511
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003512 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003513 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3514 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3515 if (index < 0 && PyErr_Occurred())
3516 return NULL; /* not an integer */
3517 /* clamp the index */
3518 if (index < 0)
3519 index = 0;
3520 else if (index > n-1)
3521 index = n-1;
3522 po->indices[i] = index;
3523 }
3524
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003525 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003526 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3527 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3528 if (index < 0 && PyErr_Occurred())
3529 return NULL; /* not an integer */
3530 if (index < 1)
3531 index = 1;
3532 else if (index > n-i)
3533 index = n-i;
3534 po->cycles[i] = index;
3535 }
3536 result = PyTuple_New(po->r);
3537 if (result == NULL)
3538 return NULL;
3539 for (i=0; i<po->r; i++) {
3540 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3541 Py_INCREF(element);
3542 PyTuple_SET_ITEM(result, i, element);
3543 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003544 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003545 Py_RETURN_NONE;
3546}
3547
3548static PyMethodDef permuations_methods[] = {
3549 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3550 reduce_doc},
3551 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3552 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003553 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3554 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003555 {NULL, NULL} /* sentinel */
3556};
3557
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003558static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003559 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003560 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003561 sizeof(permutationsobject), /* tp_basicsize */
3562 0, /* tp_itemsize */
3563 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003564 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003565 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003566 0, /* tp_getattr */
3567 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003568 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003569 0, /* tp_repr */
3570 0, /* tp_as_number */
3571 0, /* tp_as_sequence */
3572 0, /* tp_as_mapping */
3573 0, /* tp_hash */
3574 0, /* tp_call */
3575 0, /* tp_str */
3576 PyObject_GenericGetAttr, /* tp_getattro */
3577 0, /* tp_setattro */
3578 0, /* tp_as_buffer */
3579 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3580 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003581 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003582 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003583 0, /* tp_clear */
3584 0, /* tp_richcompare */
3585 0, /* tp_weaklistoffset */
3586 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003587 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003588 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003589 0, /* tp_members */
3590 0, /* tp_getset */
3591 0, /* tp_base */
3592 0, /* tp_dict */
3593 0, /* tp_descr_get */
3594 0, /* tp_descr_set */
3595 0, /* tp_dictoffset */
3596 0, /* tp_init */
3597 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003598 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003599 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003600};
3601
Tal Einatc4bccd32018-09-12 00:49:13 +03003602
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003603/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003604
3605typedef struct {
3606 PyObject_HEAD
3607 PyObject *total;
3608 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003609 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003610 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003611} accumulateobject;
3612
Tal Einatc4bccd32018-09-12 00:49:13 +03003613/*[clinic input]
3614@classmethod
3615itertools.accumulate.__new__
3616 iterable: object
3617 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003618 *
3619 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003620Return series of accumulated sums (or other binary function results).
3621[clinic start generated code]*/
3622
Raymond Hettinger482ba772010-12-01 22:48:00 +00003623static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003624itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003625 PyObject *binop, PyObject *initial)
3626/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003627{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003628 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003629 accumulateobject *lz;
3630
Raymond Hettinger482ba772010-12-01 22:48:00 +00003631 /* Get iterator. */
3632 it = PyObject_GetIter(iterable);
3633 if (it == NULL)
3634 return NULL;
3635
Raymond Hettinger482ba772010-12-01 22:48:00 +00003636 /* create accumulateobject structure */
3637 lz = (accumulateobject *)type->tp_alloc(type, 0);
3638 if (lz == NULL) {
3639 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003640 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003641 }
3642
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003643 if (binop != Py_None) {
3644 Py_XINCREF(binop);
3645 lz->binop = binop;
3646 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003647 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003648 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003649 Py_XINCREF(initial);
3650 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003651 return (PyObject *)lz;
3652}
3653
3654static void
3655accumulate_dealloc(accumulateobject *lz)
3656{
3657 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003658 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003659 Py_XDECREF(lz->total);
3660 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003661 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003662 Py_TYPE(lz)->tp_free(lz);
3663}
3664
3665static int
3666accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3667{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003668 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003669 Py_VISIT(lz->it);
3670 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003671 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003672 return 0;
3673}
3674
3675static PyObject *
3676accumulate_next(accumulateobject *lz)
3677{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003678 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003679
Lisa Roach9718b592018-09-23 17:34:59 -07003680 if (lz->initial != Py_None) {
3681 lz->total = lz->initial;
3682 Py_INCREF(Py_None);
3683 lz->initial = Py_None;
3684 Py_INCREF(lz->total);
3685 return lz->total;
3686 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003687 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003688 if (val == NULL)
3689 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003690
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003691 if (lz->total == NULL) {
3692 Py_INCREF(val);
3693 lz->total = val;
3694 return lz->total;
3695 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003696
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003697 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003698 newtotal = PyNumber_Add(lz->total, val);
3699 else
3700 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003701 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003702 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003703 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003704
Raymond Hettinger482ba772010-12-01 22:48:00 +00003705 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003706 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003707 return newtotal;
3708}
3709
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003710static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303711accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003712{
Lisa Roach9718b592018-09-23 17:34:59 -07003713 if (lz->initial != Py_None) {
3714 PyObject *it;
3715
3716 assert(lz->total == NULL);
3717 if (PyType_Ready(&chain_type) < 0)
3718 return NULL;
3719 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3720 lz->initial, lz->it);
3721 if (it == NULL)
3722 return NULL;
3723 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3724 it, lz->binop?lz->binop:Py_None, Py_None);
3725 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003726 if (lz->total == Py_None) {
3727 PyObject *it;
3728
3729 if (PyType_Ready(&chain_type) < 0)
3730 return NULL;
3731 if (PyType_Ready(&islice_type) < 0)
3732 return NULL;
3733 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3734 lz->total, lz->it);
3735 if (it == NULL)
3736 return NULL;
3737 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3738 it, lz->binop ? lz->binop : Py_None);
3739 if (it == NULL)
3740 return NULL;
3741 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3742 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003743 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3744 lz->it, lz->binop?lz->binop:Py_None,
3745 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003746}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003747
3748static PyObject *
3749accumulate_setstate(accumulateobject *lz, PyObject *state)
3750{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003751 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003752 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003753 Py_RETURN_NONE;
3754}
3755
3756static PyMethodDef accumulate_methods[] = {
3757 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3758 reduce_doc},
3759 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3760 setstate_doc},
3761 {NULL, NULL} /* sentinel */
3762};
3763
Raymond Hettinger482ba772010-12-01 22:48:00 +00003764static PyTypeObject accumulate_type = {
3765 PyVarObject_HEAD_INIT(NULL, 0)
3766 "itertools.accumulate", /* tp_name */
3767 sizeof(accumulateobject), /* tp_basicsize */
3768 0, /* tp_itemsize */
3769 /* methods */
3770 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003771 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003772 0, /* tp_getattr */
3773 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003774 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003775 0, /* tp_repr */
3776 0, /* tp_as_number */
3777 0, /* tp_as_sequence */
3778 0, /* tp_as_mapping */
3779 0, /* tp_hash */
3780 0, /* tp_call */
3781 0, /* tp_str */
3782 PyObject_GenericGetAttr, /* tp_getattro */
3783 0, /* tp_setattro */
3784 0, /* tp_as_buffer */
3785 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3786 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003787 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003788 (traverseproc)accumulate_traverse, /* tp_traverse */
3789 0, /* tp_clear */
3790 0, /* tp_richcompare */
3791 0, /* tp_weaklistoffset */
3792 PyObject_SelfIter, /* tp_iter */
3793 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003794 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003795 0, /* tp_members */
3796 0, /* tp_getset */
3797 0, /* tp_base */
3798 0, /* tp_dict */
3799 0, /* tp_descr_get */
3800 0, /* tp_descr_set */
3801 0, /* tp_dictoffset */
3802 0, /* tp_init */
3803 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003804 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003805 PyObject_GC_Del, /* tp_free */
3806};
3807
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003808
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003809/* compress object ************************************************************/
3810
3811/* Equivalent to:
3812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003813 def compress(data, selectors):
3814 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3815 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003816*/
3817
3818typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003819 PyObject_HEAD
3820 PyObject *data;
3821 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003822} compressobject;
3823
Tal Einatc4bccd32018-09-12 00:49:13 +03003824/*[clinic input]
3825@classmethod
3826itertools.compress.__new__
3827 data as seq1: object
3828 selectors as seq2: object
3829Return data elements corresponding to true selector elements.
3830
3831Forms a shorter iterator from selected data elements using the selectors to
3832choose the data elements.
3833[clinic start generated code]*/
3834
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003835static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003836itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3837/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003838{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003839 PyObject *data=NULL, *selectors=NULL;
3840 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003841
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003842 data = PyObject_GetIter(seq1);
3843 if (data == NULL)
3844 goto fail;
3845 selectors = PyObject_GetIter(seq2);
3846 if (selectors == NULL)
3847 goto fail;
3848
3849 /* create compressobject structure */
3850 lz = (compressobject *)type->tp_alloc(type, 0);
3851 if (lz == NULL)
3852 goto fail;
3853 lz->data = data;
3854 lz->selectors = selectors;
3855 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003856
3857fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003858 Py_XDECREF(data);
3859 Py_XDECREF(selectors);
3860 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003861}
3862
3863static void
3864compress_dealloc(compressobject *lz)
3865{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003866 PyObject_GC_UnTrack(lz);
3867 Py_XDECREF(lz->data);
3868 Py_XDECREF(lz->selectors);
3869 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003870}
3871
3872static int
3873compress_traverse(compressobject *lz, visitproc visit, void *arg)
3874{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003875 Py_VISIT(lz->data);
3876 Py_VISIT(lz->selectors);
3877 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003878}
3879
3880static PyObject *
3881compress_next(compressobject *lz)
3882{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003883 PyObject *data = lz->data, *selectors = lz->selectors;
3884 PyObject *datum, *selector;
3885 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3886 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3887 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003889 while (1) {
3890 /* Steps: get datum, get selector, evaluate selector.
3891 Order is important (to match the pure python version
3892 in terms of which input gets a chance to raise an
3893 exception first).
3894 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003895
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003896 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003897 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003898 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003899
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003900 selector = selectornext(selectors);
3901 if (selector == NULL) {
3902 Py_DECREF(datum);
3903 return NULL;
3904 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003905
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003906 ok = PyObject_IsTrue(selector);
3907 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003908 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003909 return datum;
3910 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003911 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003912 return NULL;
3913 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003914}
3915
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003916static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303917compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003918{
3919 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3920 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003921}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003922
3923static PyMethodDef compress_methods[] = {
3924 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3925 reduce_doc},
3926 {NULL, NULL} /* sentinel */
3927};
3928
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003929static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003930 PyVarObject_HEAD_INIT(NULL, 0)
3931 "itertools.compress", /* tp_name */
3932 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003933 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003934 /* methods */
3935 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003936 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003937 0, /* tp_getattr */
3938 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003939 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003940 0, /* tp_repr */
3941 0, /* tp_as_number */
3942 0, /* tp_as_sequence */
3943 0, /* tp_as_mapping */
3944 0, /* tp_hash */
3945 0, /* tp_call */
3946 0, /* tp_str */
3947 PyObject_GenericGetAttr, /* tp_getattro */
3948 0, /* tp_setattro */
3949 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003950 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003951 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003952 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003953 (traverseproc)compress_traverse, /* tp_traverse */
3954 0, /* tp_clear */
3955 0, /* tp_richcompare */
3956 0, /* tp_weaklistoffset */
3957 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003958 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003959 compress_methods, /* tp_methods */
3960 0, /* tp_members */
3961 0, /* tp_getset */
3962 0, /* tp_base */
3963 0, /* tp_dict */
3964 0, /* tp_descr_get */
3965 0, /* tp_descr_set */
3966 0, /* tp_dictoffset */
3967 0, /* tp_init */
3968 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003969 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003970 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003971};
3972
3973
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003974/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003975
3976typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003977 PyObject_HEAD
3978 PyObject *func;
3979 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003980} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003981
Tal Einatc4bccd32018-09-12 00:49:13 +03003982/*[clinic input]
3983@classmethod
3984itertools.filterfalse.__new__
3985 function as func: object
3986 iterable as seq: object
3987 /
3988Return those items of iterable for which function(item) is false.
3989
3990If function is None, return the items that are false.
3991[clinic start generated code]*/
3992
Raymond Hettinger60eca932003-02-09 06:40:58 +00003993static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003994itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3995/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003996{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003997 PyObject *it;
3998 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003999
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004000 /* Get iterator. */
4001 it = PyObject_GetIter(seq);
4002 if (it == NULL)
4003 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004005 /* create filterfalseobject structure */
4006 lz = (filterfalseobject *)type->tp_alloc(type, 0);
4007 if (lz == NULL) {
4008 Py_DECREF(it);
4009 return NULL;
4010 }
4011 Py_INCREF(func);
4012 lz->func = func;
4013 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004014
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004015 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004016}
4017
4018static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004019filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00004020{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004021 PyObject_GC_UnTrack(lz);
4022 Py_XDECREF(lz->func);
4023 Py_XDECREF(lz->it);
4024 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00004025}
4026
4027static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004028filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00004029{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004030 Py_VISIT(lz->it);
4031 Py_VISIT(lz->func);
4032 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004033}
4034
4035static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004036filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00004037{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004038 PyObject *item;
4039 PyObject *it = lz->it;
4040 long ok;
4041 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00004042
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 iternext = *Py_TYPE(it)->tp_iternext;
4044 for (;;) {
4045 item = iternext(it);
4046 if (item == NULL)
4047 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004048
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004049 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
4050 ok = PyObject_IsTrue(item);
4051 } else {
4052 PyObject *good;
Petr Viktorinffd97532020-02-11 17:46:57 +01004053 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004054 if (good == NULL) {
4055 Py_DECREF(item);
4056 return NULL;
4057 }
4058 ok = PyObject_IsTrue(good);
4059 Py_DECREF(good);
4060 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02004061 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004062 return item;
4063 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02004064 if (ok < 0)
4065 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004066 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00004067}
4068
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004069static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304070filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004071{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004072 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
4073}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004074
4075static PyMethodDef filterfalse_methods[] = {
4076 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
4077 reduce_doc},
4078 {NULL, NULL} /* sentinel */
4079};
4080
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004081static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004082 PyVarObject_HEAD_INIT(NULL, 0)
4083 "itertools.filterfalse", /* tp_name */
4084 sizeof(filterfalseobject), /* tp_basicsize */
4085 0, /* tp_itemsize */
4086 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004087 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004088 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004089 0, /* tp_getattr */
4090 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004091 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004092 0, /* tp_repr */
4093 0, /* tp_as_number */
4094 0, /* tp_as_sequence */
4095 0, /* tp_as_mapping */
4096 0, /* tp_hash */
4097 0, /* tp_call */
4098 0, /* tp_str */
4099 PyObject_GenericGetAttr, /* tp_getattro */
4100 0, /* tp_setattro */
4101 0, /* tp_as_buffer */
4102 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4103 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004104 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004105 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004106 0, /* tp_clear */
4107 0, /* tp_richcompare */
4108 0, /* tp_weaklistoffset */
4109 PyObject_SelfIter, /* tp_iter */
4110 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004111 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004112 0, /* tp_members */
4113 0, /* tp_getset */
4114 0, /* tp_base */
4115 0, /* tp_dict */
4116 0, /* tp_descr_get */
4117 0, /* tp_descr_set */
4118 0, /* tp_dictoffset */
4119 0, /* tp_init */
4120 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004121 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004122 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004123};
4124
4125
4126/* count object ************************************************************/
4127
4128typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004129 PyObject_HEAD
4130 Py_ssize_t cnt;
4131 PyObject *long_cnt;
4132 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004133} countobject;
4134
Raymond Hettinger30729212009-02-12 06:28:27 +00004135/* Counting logic and invariants:
4136
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004137fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00004138
Serhiy Storchaka483405b2015-02-17 10:14:30 +02004139 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004140 Advances with: cnt += 1
4141 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00004142
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004143slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004144
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004145 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4146 All counting is done with python objects (no overflows or underflows).
4147 Advances with: long_cnt += long_step
4148 Step may be zero -- effectively a slow version of repeat(cnt).
4149 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004150*/
4151
Tal Einatc4bccd32018-09-12 00:49:13 +03004152/*[clinic input]
4153@classmethod
4154itertools.count.__new__
4155 start as long_cnt: object(c_default="NULL") = 0
4156 step as long_step: object(c_default="NULL") = 1
4157Return a count object whose .__next__() method returns consecutive values.
4158
4159Equivalent to:
4160 def count(firstval=0, step=1):
4161 x = firstval
4162 while 1:
4163 yield x
4164 x += step
4165[clinic start generated code]*/
4166
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004167static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004168itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4169 PyObject *long_step)
4170/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004171{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004172 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004173 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004174 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004175 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004176
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004177 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4178 (long_step != NULL && !PyNumber_Check(long_step))) {
4179 PyErr_SetString(PyExc_TypeError, "a number is required");
4180 return NULL;
4181 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004182
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004183 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4184 (long_step == NULL || PyLong_Check(long_step));
4185
4186 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004187 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004188 if (fast_mode) {
4189 assert(PyLong_Check(long_cnt));
4190 cnt = PyLong_AsSsize_t(long_cnt);
4191 if (cnt == -1 && PyErr_Occurred()) {
4192 PyErr_Clear();
4193 fast_mode = 0;
4194 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004195 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 } else {
4197 cnt = 0;
Victor Stinner37834132020-10-27 17:12:53 +01004198 long_cnt = _PyLong_GetZero();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004199 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004200 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004201
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004202 /* If not specified, step defaults to 1 */
Victor Stinner37834132020-10-27 17:12:53 +01004203 if (long_step == NULL) {
4204 long_step = _PyLong_GetOne();
4205 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004206 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004207
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004208 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004209
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004210 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004211 if (fast_mode) {
4212 assert(PyLong_Check(long_step));
4213 step = PyLong_AsLong(long_step);
4214 if (step != 1) {
4215 fast_mode = 0;
4216 if (step == -1 && PyErr_Occurred())
4217 PyErr_Clear();
4218 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004219 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004220
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004221 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004222 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004223 else
4224 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004225
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004226 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4227 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4228 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004229 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004231 /* create countobject structure */
4232 lz = (countobject *)type->tp_alloc(type, 0);
4233 if (lz == NULL) {
4234 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004235 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004236 return NULL;
4237 }
4238 lz->cnt = cnt;
4239 lz->long_cnt = long_cnt;
4240 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004242 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004243}
4244
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004245static void
4246count_dealloc(countobject *lz)
4247{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004248 PyObject_GC_UnTrack(lz);
4249 Py_XDECREF(lz->long_cnt);
4250 Py_XDECREF(lz->long_step);
4251 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004252}
4253
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004254static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004255count_traverse(countobject *lz, visitproc visit, void *arg)
4256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004257 Py_VISIT(lz->long_cnt);
4258 Py_VISIT(lz->long_step);
4259 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004260}
4261
4262static PyObject *
4263count_nextlong(countobject *lz)
4264{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004265 PyObject *long_cnt;
4266 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004267
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004268 long_cnt = lz->long_cnt;
4269 if (long_cnt == NULL) {
4270 /* Switch to slow_mode */
4271 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4272 if (long_cnt == NULL)
4273 return NULL;
4274 }
4275 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004276
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004277 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4278 if (stepped_up == NULL)
4279 return NULL;
4280 lz->long_cnt = stepped_up;
4281 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004282}
4283
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004284static PyObject *
4285count_next(countobject *lz)
4286{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004287 if (lz->cnt == PY_SSIZE_T_MAX)
4288 return count_nextlong(lz);
4289 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004290}
4291
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004292static PyObject *
4293count_repr(countobject *lz)
4294{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004295 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004296 return PyUnicode_FromFormat("%s(%zd)",
4297 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004298
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004299 if (PyLong_Check(lz->long_step)) {
4300 long step = PyLong_AsLong(lz->long_step);
4301 if (step == -1 && PyErr_Occurred()) {
4302 PyErr_Clear();
4303 }
4304 if (step == 1) {
4305 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004306 return PyUnicode_FromFormat("%s(%R)",
4307 _PyType_Name(Py_TYPE(lz)),
4308 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004309 }
4310 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004311 return PyUnicode_FromFormat("%s(%R, %R)",
4312 _PyType_Name(Py_TYPE(lz)),
4313 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004314}
4315
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004316static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304317count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004318{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004319 if (lz->cnt == PY_SSIZE_T_MAX)
4320 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4321 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004322}
4323
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004324static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004325 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004326 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004327 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004328};
4329
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004330static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004331 PyVarObject_HEAD_INIT(NULL, 0)
4332 "itertools.count", /* tp_name */
4333 sizeof(countobject), /* tp_basicsize */
4334 0, /* tp_itemsize */
4335 /* methods */
4336 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004337 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004338 0, /* tp_getattr */
4339 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004340 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004341 (reprfunc)count_repr, /* tp_repr */
4342 0, /* tp_as_number */
4343 0, /* tp_as_sequence */
4344 0, /* tp_as_mapping */
4345 0, /* tp_hash */
4346 0, /* tp_call */
4347 0, /* tp_str */
4348 PyObject_GenericGetAttr, /* tp_getattro */
4349 0, /* tp_setattro */
4350 0, /* tp_as_buffer */
4351 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004352 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004353 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004354 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004355 0, /* tp_clear */
4356 0, /* tp_richcompare */
4357 0, /* tp_weaklistoffset */
4358 PyObject_SelfIter, /* tp_iter */
4359 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004360 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004361 0, /* tp_members */
4362 0, /* tp_getset */
4363 0, /* tp_base */
4364 0, /* tp_dict */
4365 0, /* tp_descr_get */
4366 0, /* tp_descr_set */
4367 0, /* tp_dictoffset */
4368 0, /* tp_init */
4369 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004370 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004371 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004372};
4373
4374
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004375/* repeat object ************************************************************/
4376
4377typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004378 PyObject_HEAD
4379 PyObject *element;
4380 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004381} repeatobject;
4382
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004383static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004384
4385static PyObject *
4386repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4387{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004388 repeatobject *ro;
4389 PyObject *element;
Raymond Hettingeraf464502019-11-09 20:28:31 -08004390 Py_ssize_t cnt = -1, n_args;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004391 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004392
Raymond Hettingeraf464502019-11-09 20:28:31 -08004393 n_args = PyTuple_GET_SIZE(args);
4394 if (kwds != NULL)
4395 n_args += PyDict_GET_SIZE(kwds);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004396 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4397 &element, &cnt))
4398 return NULL;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004399 /* Does user supply times argument? */
Raymond Hettingeraf464502019-11-09 20:28:31 -08004400 if (n_args == 2 && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004401 cnt = 0;
4402
4403 ro = (repeatobject *)type->tp_alloc(type, 0);
4404 if (ro == NULL)
4405 return NULL;
4406 Py_INCREF(element);
4407 ro->element = element;
4408 ro->cnt = cnt;
4409 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004410}
4411
4412static void
4413repeat_dealloc(repeatobject *ro)
4414{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004415 PyObject_GC_UnTrack(ro);
4416 Py_XDECREF(ro->element);
4417 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004418}
4419
4420static int
4421repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4422{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004423 Py_VISIT(ro->element);
4424 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004425}
4426
4427static PyObject *
4428repeat_next(repeatobject *ro)
4429{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004430 if (ro->cnt == 0)
4431 return NULL;
4432 if (ro->cnt > 0)
4433 ro->cnt--;
4434 Py_INCREF(ro->element);
4435 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004436}
4437
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004438static PyObject *
4439repeat_repr(repeatobject *ro)
4440{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004441 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004442 return PyUnicode_FromFormat("%s(%R)",
4443 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004444 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004445 return PyUnicode_FromFormat("%s(%R, %zd)",
4446 _PyType_Name(Py_TYPE(ro)), ro->element,
4447 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004448}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004449
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004450static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304451repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004452{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004453 if (ro->cnt == -1) {
4454 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4455 return NULL;
4456 }
4457 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004458}
4459
Armin Rigof5b3e362006-02-11 21:32:43 +00004460PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004461
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004462static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304463repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004464{
4465 /* unpickle this so that a new repeat iterator is constructed with an
4466 * object, then call __setstate__ on it to set cnt
4467 */
4468 if (ro->cnt >= 0)
4469 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4470 else
4471 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4472}
4473
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004474static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004475 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004476 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004477 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004478};
4479
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004480PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004481"repeat(object [,times]) -> create an iterator which returns the object\n\
4482for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004483endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004484
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004485static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004486 PyVarObject_HEAD_INIT(NULL, 0)
4487 "itertools.repeat", /* tp_name */
4488 sizeof(repeatobject), /* tp_basicsize */
4489 0, /* tp_itemsize */
4490 /* methods */
4491 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004492 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004493 0, /* tp_getattr */
4494 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004495 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004496 (reprfunc)repeat_repr, /* tp_repr */
4497 0, /* tp_as_number */
4498 0, /* tp_as_sequence */
4499 0, /* tp_as_mapping */
4500 0, /* tp_hash */
4501 0, /* tp_call */
4502 0, /* tp_str */
4503 PyObject_GenericGetAttr, /* tp_getattro */
4504 0, /* tp_setattro */
4505 0, /* tp_as_buffer */
4506 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4507 Py_TPFLAGS_BASETYPE, /* tp_flags */
4508 repeat_doc, /* tp_doc */
4509 (traverseproc)repeat_traverse, /* tp_traverse */
4510 0, /* tp_clear */
4511 0, /* tp_richcompare */
4512 0, /* tp_weaklistoffset */
4513 PyObject_SelfIter, /* tp_iter */
4514 (iternextfunc)repeat_next, /* tp_iternext */
4515 repeat_methods, /* tp_methods */
4516 0, /* tp_members */
4517 0, /* tp_getset */
4518 0, /* tp_base */
4519 0, /* tp_dict */
4520 0, /* tp_descr_get */
4521 0, /* tp_descr_set */
4522 0, /* tp_dictoffset */
4523 0, /* tp_init */
4524 0, /* tp_alloc */
4525 repeat_new, /* tp_new */
4526 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004527};
4528
Tal Einatc4bccd32018-09-12 00:49:13 +03004529
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004530/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004531
4532typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004533 PyObject_HEAD
4534 Py_ssize_t tuplesize;
4535 Py_ssize_t numactive;
4536 PyObject *ittuple; /* tuple of iterators */
4537 PyObject *result;
4538 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004539} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004540
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004541static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004542
4543static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004544zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004545{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004546 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004547 ziplongestobject *lz;
4548 Py_ssize_t i;
4549 PyObject *ittuple; /* tuple of iterators */
4550 PyObject *result;
4551 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004552 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004553
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004554 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004555 fillvalue = NULL;
4556 if (PyDict_GET_SIZE(kwds) == 1) {
4557 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4558 }
4559 if (fillvalue == NULL) {
4560 if (!PyErr_Occurred()) {
4561 PyErr_SetString(PyExc_TypeError,
4562 "zip_longest() got an unexpected keyword argument");
4563 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004564 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004565 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004566 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004567
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004568 /* args must be a tuple */
4569 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004570 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004571
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004572 /* obtain iterators */
4573 ittuple = PyTuple_New(tuplesize);
4574 if (ittuple == NULL)
4575 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004576 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004577 PyObject *item = PyTuple_GET_ITEM(args, i);
4578 PyObject *it = PyObject_GetIter(item);
4579 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004580 Py_DECREF(ittuple);
4581 return NULL;
4582 }
4583 PyTuple_SET_ITEM(ittuple, i, it);
4584 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004585
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004586 /* create a result holder */
4587 result = PyTuple_New(tuplesize);
4588 if (result == NULL) {
4589 Py_DECREF(ittuple);
4590 return NULL;
4591 }
4592 for (i=0 ; i < tuplesize ; i++) {
4593 Py_INCREF(Py_None);
4594 PyTuple_SET_ITEM(result, i, Py_None);
4595 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004596
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004597 /* create ziplongestobject structure */
4598 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4599 if (lz == NULL) {
4600 Py_DECREF(ittuple);
4601 Py_DECREF(result);
4602 return NULL;
4603 }
4604 lz->ittuple = ittuple;
4605 lz->tuplesize = tuplesize;
4606 lz->numactive = tuplesize;
4607 lz->result = result;
4608 Py_INCREF(fillvalue);
4609 lz->fillvalue = fillvalue;
4610 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004611}
4612
4613static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004614zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004615{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004616 PyObject_GC_UnTrack(lz);
4617 Py_XDECREF(lz->ittuple);
4618 Py_XDECREF(lz->result);
4619 Py_XDECREF(lz->fillvalue);
4620 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004621}
4622
4623static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004624zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004625{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004626 Py_VISIT(lz->ittuple);
4627 Py_VISIT(lz->result);
4628 Py_VISIT(lz->fillvalue);
4629 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004630}
4631
4632static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004633zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004634{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004635 Py_ssize_t i;
4636 Py_ssize_t tuplesize = lz->tuplesize;
4637 PyObject *result = lz->result;
4638 PyObject *it;
4639 PyObject *item;
4640 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004641
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004642 if (tuplesize == 0)
4643 return NULL;
4644 if (lz->numactive == 0)
4645 return NULL;
4646 if (Py_REFCNT(result) == 1) {
4647 Py_INCREF(result);
4648 for (i=0 ; i < tuplesize ; i++) {
4649 it = PyTuple_GET_ITEM(lz->ittuple, i);
4650 if (it == NULL) {
4651 Py_INCREF(lz->fillvalue);
4652 item = lz->fillvalue;
4653 } else {
4654 item = PyIter_Next(it);
4655 if (item == NULL) {
4656 lz->numactive -= 1;
4657 if (lz->numactive == 0 || PyErr_Occurred()) {
4658 lz->numactive = 0;
4659 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004660 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004661 } else {
4662 Py_INCREF(lz->fillvalue);
4663 item = lz->fillvalue;
4664 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4665 Py_DECREF(it);
4666 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004667 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004668 }
4669 olditem = PyTuple_GET_ITEM(result, i);
4670 PyTuple_SET_ITEM(result, i, item);
4671 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004672 }
Brandt Bucher226a0122020-12-04 19:45:57 -08004673 // bpo-42536: The GC may have untracked this result tuple. Since we're
4674 // recycling it, make sure it's tracked again:
4675 if (!_PyObject_GC_IS_TRACKED(result)) {
4676 _PyObject_GC_TRACK(result);
4677 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004678 } else {
4679 result = PyTuple_New(tuplesize);
4680 if (result == NULL)
4681 return NULL;
4682 for (i=0 ; i < tuplesize ; i++) {
4683 it = PyTuple_GET_ITEM(lz->ittuple, i);
4684 if (it == NULL) {
4685 Py_INCREF(lz->fillvalue);
4686 item = lz->fillvalue;
4687 } else {
4688 item = PyIter_Next(it);
4689 if (item == NULL) {
4690 lz->numactive -= 1;
4691 if (lz->numactive == 0 || PyErr_Occurred()) {
4692 lz->numactive = 0;
4693 Py_DECREF(result);
4694 return NULL;
4695 } else {
4696 Py_INCREF(lz->fillvalue);
4697 item = lz->fillvalue;
4698 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4699 Py_DECREF(it);
4700 }
4701 }
4702 }
4703 PyTuple_SET_ITEM(result, i, item);
4704 }
4705 }
4706 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004707}
4708
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004709static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304710zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004711{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004712
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004713 /* Create a new tuple with empty sequences where appropriate to pickle.
4714 * Then use setstate to set the fillvalue
4715 */
4716 int i;
4717 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004718
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004719 if (args == NULL)
4720 return NULL;
4721 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4722 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4723 if (elem == NULL) {
4724 elem = PyTuple_New(0);
4725 if (elem == NULL) {
4726 Py_DECREF(args);
4727 return NULL;
4728 }
4729 } else
4730 Py_INCREF(elem);
4731 PyTuple_SET_ITEM(args, i, elem);
4732 }
4733 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4734}
4735
4736static PyObject *
4737zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4738{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004739 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004740 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004741 Py_RETURN_NONE;
4742}
4743
4744static PyMethodDef zip_longest_methods[] = {
4745 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4746 reduce_doc},
4747 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4748 setstate_doc},
4749 {NULL, NULL} /* sentinel */
4750};
4751
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004752PyDoc_STRVAR(zip_longest_doc,
4753"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004754\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004755Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004756the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004757method continues until the longest iterable in the argument sequence\n\
4758is exhausted and then it raises StopIteration. When the shorter iterables\n\
4759are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4760defaults to None or can be specified by a keyword argument.\n\
4761");
4762
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004763static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004764 PyVarObject_HEAD_INIT(NULL, 0)
4765 "itertools.zip_longest", /* tp_name */
4766 sizeof(ziplongestobject), /* tp_basicsize */
4767 0, /* tp_itemsize */
4768 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004769 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004770 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004771 0, /* tp_getattr */
4772 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004773 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004774 0, /* tp_repr */
4775 0, /* tp_as_number */
4776 0, /* tp_as_sequence */
4777 0, /* tp_as_mapping */
4778 0, /* tp_hash */
4779 0, /* tp_call */
4780 0, /* tp_str */
4781 PyObject_GenericGetAttr, /* tp_getattro */
4782 0, /* tp_setattro */
4783 0, /* tp_as_buffer */
4784 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4785 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004786 zip_longest_doc, /* tp_doc */
4787 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004788 0, /* tp_clear */
4789 0, /* tp_richcompare */
4790 0, /* tp_weaklistoffset */
4791 PyObject_SelfIter, /* tp_iter */
4792 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004793 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004794 0, /* tp_members */
4795 0, /* tp_getset */
4796 0, /* tp_base */
4797 0, /* tp_dict */
4798 0, /* tp_descr_get */
4799 0, /* tp_descr_set */
4800 0, /* tp_dictoffset */
4801 0, /* tp_init */
4802 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004803 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004804 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004805};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004806
Tal Einatc4bccd32018-09-12 00:49:13 +03004807
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004808/* module level code ********************************************************/
4809
4810PyDoc_STRVAR(module_doc,
4811"Functional tools for creating and using iterators.\n\
4812\n\
4813Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004814count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004815cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004816repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004817\n\
4818Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004819accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004820chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4821chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004822compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4823dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4824groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004825filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004826islice(seq, [start,] stop [, step]) --> elements from\n\
4827 seq[start:stop:step]\n\
Raymond Hettingercc061d02020-11-30 20:42:54 -08004828pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004829starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004830tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004831takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004832zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004833\n\
4834Combinatoric generators:\n\
4835product(p, q, ... [repeat=1]) --> cartesian product\n\
4836permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004837combinations(p, r)\n\
4838combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004839");
4840
Dong-hee Na514c4692020-03-18 02:46:24 +09004841static int
4842itertoolsmodule_exec(PyObject *m)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004843{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004844 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004845 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004846 &combinations_type,
4847 &cwr_type,
4848 &cycle_type,
4849 &dropwhile_type,
4850 &takewhile_type,
4851 &islice_type,
4852 &starmap_type,
4853 &chain_type,
4854 &compress_type,
4855 &filterfalse_type,
4856 &count_type,
4857 &ziplongest_type,
Raymond Hettingercc061d02020-11-30 20:42:54 -08004858 &pairwise_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004859 &permutations_type,
4860 &product_type,
4861 &repeat_type,
4862 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004863 &_grouper_type,
4864 &tee_type,
Dong-hee Na05e4a292020-03-23 01:17:34 +09004865 &teedataobject_type
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004866 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004867
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004868 Py_SET_TYPE(&teedataobject_type, &PyType_Type);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004869
Dong-hee Na05e4a292020-03-23 01:17:34 +09004870 for (size_t i = 0; i < Py_ARRAY_LENGTH(typelist); i++) {
4871 if (PyModule_AddType(m, typelist[i]) < 0) {
Dong-hee Na514c4692020-03-18 02:46:24 +09004872 return -1;
4873 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004874 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004875
Dong-hee Na514c4692020-03-18 02:46:24 +09004876 return 0;
4877}
4878
4879static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4880 {Py_mod_exec, itertoolsmodule_exec},
4881 {0, NULL}
4882};
4883
4884static PyMethodDef module_methods[] = {
4885 ITERTOOLS_TEE_METHODDEF
4886 {NULL, NULL} /* sentinel */
4887};
4888
4889
4890static struct PyModuleDef itertoolsmodule = {
4891 PyModuleDef_HEAD_INIT,
4892 "itertools",
4893 module_doc,
4894 0,
4895 module_methods,
4896 itertoolsmodule_slots,
4897 NULL,
4898 NULL,
4899 NULL
4900};
4901
4902PyMODINIT_FUNC
4903PyInit_itertools(void)
4904{
4905 return PyModuleDef_Init(&itertoolsmodule);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004906}