blob: ec8f0ae14206cbbaf8f607f3bc3dcd1962aabe85 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
Raymond Hettingerca3788c2015-08-16 14:49:24 -07002#define PY_SSIZE_T_CLEAN
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00004#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00006/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008*/
9
Tal Einat3286ce42018-09-10 21:33:08 +030010/*[clinic input]
11module itertools
12class itertools.groupby "groupbyobject *" "&groupby_type"
13class itertools._grouper "_grouperobject *" "&_grouper_type"
Tal Einatc4bccd32018-09-12 00:49:13 +030014class itertools.teedataobject "teedataobject *" "&teedataobject_type"
15class itertools._tee "teeobject *" "&tee_type"
16class itertools.cycle "cycleobject *" "&cycle_type"
17class itertools.dropwhile "dropwhileobject *" "&dropwhile_type"
18class itertools.takewhile "takewhileobject *" "&takewhile_type"
19class itertools.starmap "starmapobject *" "&starmap_type"
20class itertools.chain "chainobject *" "&chain_type"
21class itertools.combinations "combinationsobject *" "&combinations_type"
22class itertools.combinations_with_replacement "cwr_object *" "&cwr_type"
23class itertools.permutations "permutationsobject *" "&permutations_type"
24class itertools.accumulate "accumulateobject *" "&accumulate_type"
25class itertools.compress "compressobject *" "&compress_type"
26class itertools.filterfalse "filterfalseobject *" "&filterfalse_type"
27class itertools.count "countobject *" "&count_type"
Tal Einat3286ce42018-09-10 21:33:08 +030028[clinic start generated code]*/
Tal Einatc4bccd32018-09-12 00:49:13 +030029/*[clinic end generated code: output=da39a3ee5e6b4b0d input=ea05c93c6d94726a]*/
Tal Einat3286ce42018-09-10 21:33:08 +030030
31static PyTypeObject groupby_type;
32static PyTypeObject _grouper_type;
Tal Einatc4bccd32018-09-12 00:49:13 +030033static PyTypeObject teedataobject_type;
34static PyTypeObject tee_type;
35static PyTypeObject cycle_type;
36static PyTypeObject dropwhile_type;
37static PyTypeObject takewhile_type;
38static PyTypeObject starmap_type;
39static PyTypeObject combinations_type;
40static PyTypeObject cwr_type;
41static PyTypeObject permutations_type;
42static PyTypeObject accumulate_type;
43static PyTypeObject compress_type;
44static PyTypeObject filterfalse_type;
45static PyTypeObject count_type;
46
Tal Einat3286ce42018-09-10 21:33:08 +030047#include "clinic/itertoolsmodule.c.h"
48
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000049
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070050/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051
52typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000053 PyObject_HEAD
54 PyObject *it;
55 PyObject *keyfunc;
56 PyObject *tgtkey;
57 PyObject *currkey;
58 PyObject *currvalue;
Serhiy Storchakac247caf2017-09-24 13:36:11 +030059 const void *currgrouper; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000060} groupbyobject;
61
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000062static PyObject *_grouper_create(groupbyobject *, PyObject *);
63
Tal Einat3286ce42018-09-10 21:33:08 +030064/*[clinic input]
65@classmethod
66itertools.groupby.__new__
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000067
Tal Einat3286ce42018-09-10 21:33:08 +030068 iterable as it: object
69 Elements to divide into groups according to the key function.
70 key as keyfunc: object = None
71 A function for computing the group category for each element.
72 If the key function is not specified or is None, the element itself
73 is used for grouping.
74
75make an iterator that returns consecutive keys and groups from the iterable
76[clinic start generated code]*/
77
78static PyObject *
79itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
80/*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
81{
82 groupbyobject *gbo;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000083
84 gbo = (groupbyobject *)type->tp_alloc(type, 0);
85 if (gbo == NULL)
86 return NULL;
87 gbo->tgtkey = NULL;
88 gbo->currkey = NULL;
89 gbo->currvalue = NULL;
90 gbo->keyfunc = keyfunc;
91 Py_INCREF(keyfunc);
92 gbo->it = PyObject_GetIter(it);
93 if (gbo->it == NULL) {
94 Py_DECREF(gbo);
95 return NULL;
96 }
97 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000098}
99
100static void
101groupby_dealloc(groupbyobject *gbo)
102{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000103 PyObject_GC_UnTrack(gbo);
104 Py_XDECREF(gbo->it);
105 Py_XDECREF(gbo->keyfunc);
106 Py_XDECREF(gbo->tgtkey);
107 Py_XDECREF(gbo->currkey);
108 Py_XDECREF(gbo->currvalue);
109 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000110}
111
112static int
113groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
114{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000115 Py_VISIT(gbo->it);
116 Py_VISIT(gbo->keyfunc);
117 Py_VISIT(gbo->tgtkey);
118 Py_VISIT(gbo->currkey);
119 Py_VISIT(gbo->currvalue);
120 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121}
122
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300123Py_LOCAL_INLINE(int)
124groupby_step(groupbyobject *gbo)
125{
126 PyObject *newvalue, *newkey, *oldvalue;
127
128 newvalue = PyIter_Next(gbo->it);
129 if (newvalue == NULL)
130 return -1;
131
132 if (gbo->keyfunc == Py_None) {
133 newkey = newvalue;
134 Py_INCREF(newvalue);
135 } else {
136 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
137 if (newkey == NULL) {
138 Py_DECREF(newvalue);
139 return -1;
140 }
141 }
142
143 oldvalue = gbo->currvalue;
144 gbo->currvalue = newvalue;
145 Py_XSETREF(gbo->currkey, newkey);
146 Py_XDECREF(oldvalue);
147 return 0;
148}
149
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000150static PyObject *
151groupby_next(groupbyobject *gbo)
152{
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300153 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000154
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300155 gbo->currgrouper = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000156 /* skip to next iteration group */
157 for (;;) {
158 if (gbo->currkey == NULL)
159 /* pass */;
160 else if (gbo->tgtkey == NULL)
161 break;
162 else {
163 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000164
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700165 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000166 if (rcmp == -1)
167 return NULL;
168 else if (rcmp == 0)
169 break;
170 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000171
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300172 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000173 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000174 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000175 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300176 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000178 grouper = _grouper_create(gbo, gbo->tgtkey);
179 if (grouper == NULL)
180 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000181
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000182 r = PyTuple_Pack(2, gbo->currkey, grouper);
183 Py_DECREF(grouper);
184 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000185}
186
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000187static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530188groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000189{
190 /* reduce as a 'new' call with an optional 'setstate' if groupby
191 * has started
192 */
193 PyObject *value;
194 if (lz->tgtkey && lz->currkey && lz->currvalue)
195 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
196 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
197 else
198 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
199 lz->it, lz->keyfunc);
200
201 return value;
202}
203
204PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
205
206static PyObject *
207groupby_setstate(groupbyobject *lz, PyObject *state)
208{
209 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300210 if (!PyTuple_Check(state)) {
211 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000212 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300213 }
214 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
215 return NULL;
216 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200217 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300218 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200219 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300220 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200221 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300222 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000223 Py_RETURN_NONE;
224}
225
226PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
227
228static PyMethodDef groupby_methods[] = {
229 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
230 reduce_doc},
231 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
232 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700233 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000234};
235
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000237 PyVarObject_HEAD_INIT(NULL, 0)
238 "itertools.groupby", /* tp_name */
239 sizeof(groupbyobject), /* tp_basicsize */
240 0, /* tp_itemsize */
241 /* methods */
242 (destructor)groupby_dealloc, /* tp_dealloc */
243 0, /* tp_print */
244 0, /* tp_getattr */
245 0, /* tp_setattr */
246 0, /* tp_reserved */
247 0, /* tp_repr */
248 0, /* tp_as_number */
249 0, /* tp_as_sequence */
250 0, /* tp_as_mapping */
251 0, /* tp_hash */
252 0, /* tp_call */
253 0, /* tp_str */
254 PyObject_GenericGetAttr, /* tp_getattro */
255 0, /* tp_setattro */
256 0, /* tp_as_buffer */
257 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
258 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einat3286ce42018-09-10 21:33:08 +0300259 itertools_groupby__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000260 (traverseproc)groupby_traverse, /* tp_traverse */
261 0, /* tp_clear */
262 0, /* tp_richcompare */
263 0, /* tp_weaklistoffset */
264 PyObject_SelfIter, /* tp_iter */
265 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000266 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000267 0, /* tp_members */
268 0, /* tp_getset */
269 0, /* tp_base */
270 0, /* tp_dict */
271 0, /* tp_descr_get */
272 0, /* tp_descr_set */
273 0, /* tp_dictoffset */
274 0, /* tp_init */
275 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300276 itertools_groupby, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000277 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000278};
279
280
281/* _grouper object (internal) ************************************************/
282
283typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000284 PyObject_HEAD
285 PyObject *parent;
286 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000287} _grouperobject;
288
Tal Einat3286ce42018-09-10 21:33:08 +0300289/*[clinic input]
290@classmethod
291itertools._grouper.__new__
292
293 parent: object(subclass_of='&groupby_type')
294 tgtkey: object
295 /
296[clinic start generated code]*/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000297
298static PyObject *
Tal Einat3286ce42018-09-10 21:33:08 +0300299itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
300 PyObject *tgtkey)
301/*[clinic end generated code: output=462efb1cdebb5914 input=dc180d7771fc8c59]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000302{
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000303 return _grouper_create((groupbyobject*) parent, tgtkey);
304}
305
306static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000307_grouper_create(groupbyobject *parent, PyObject *tgtkey)
308{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000309 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000310
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000311 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
312 if (igo == NULL)
313 return NULL;
314 igo->parent = (PyObject *)parent;
315 Py_INCREF(parent);
316 igo->tgtkey = tgtkey;
317 Py_INCREF(tgtkey);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300318 parent->currgrouper = igo; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000319
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000320 PyObject_GC_Track(igo);
321 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000322}
323
324static void
325_grouper_dealloc(_grouperobject *igo)
326{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000327 PyObject_GC_UnTrack(igo);
328 Py_DECREF(igo->parent);
329 Py_DECREF(igo->tgtkey);
330 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000331}
332
333static int
334_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
335{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000336 Py_VISIT(igo->parent);
337 Py_VISIT(igo->tgtkey);
338 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000339}
340
341static PyObject *
342_grouper_next(_grouperobject *igo)
343{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000344 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300345 PyObject *r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000346 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000347
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300348 if (gbo->currgrouper != igo)
349 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000350 if (gbo->currvalue == NULL) {
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300351 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000352 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000353 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000354
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000355 assert(gbo->currkey != NULL);
356 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
357 if (rcmp <= 0)
358 /* got any error or current group is end */
359 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000360
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000361 r = gbo->currvalue;
362 gbo->currvalue = NULL;
363 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000364
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000365 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000366}
367
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000368static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530369_grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000370{
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300371 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
372 return Py_BuildValue("N(())", _PyObject_GetBuiltin("iter"));
373 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700374 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000375}
376
377static PyMethodDef _grouper_methods[] = {
378 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
379 reduce_doc},
380 {NULL, NULL} /* sentinel */
381};
382
383
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000384static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000385 PyVarObject_HEAD_INIT(NULL, 0)
386 "itertools._grouper", /* tp_name */
387 sizeof(_grouperobject), /* tp_basicsize */
388 0, /* tp_itemsize */
389 /* methods */
390 (destructor)_grouper_dealloc, /* tp_dealloc */
391 0, /* tp_print */
392 0, /* tp_getattr */
393 0, /* tp_setattr */
394 0, /* tp_reserved */
395 0, /* tp_repr */
396 0, /* tp_as_number */
397 0, /* tp_as_sequence */
398 0, /* tp_as_mapping */
399 0, /* tp_hash */
400 0, /* tp_call */
401 0, /* tp_str */
402 PyObject_GenericGetAttr, /* tp_getattro */
403 0, /* tp_setattro */
404 0, /* tp_as_buffer */
405 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
406 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700407 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000408 0, /* tp_clear */
409 0, /* tp_richcompare */
410 0, /* tp_weaklistoffset */
411 PyObject_SelfIter, /* tp_iter */
412 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000413 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000414 0, /* tp_members */
415 0, /* tp_getset */
416 0, /* tp_base */
417 0, /* tp_dict */
418 0, /* tp_descr_get */
419 0, /* tp_descr_set */
420 0, /* tp_dictoffset */
421 0, /* tp_init */
422 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300423 itertools__grouper, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000424 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000425};
426
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700428/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000429
Raymond Hettingerad983e72003-11-12 14:32:26 +0000430/* The teedataobject pre-allocates space for LINKCELLS number of objects.
431 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000432 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433 the other structure members including PyHEAD overhead. The larger the
434 value, the less memory overhead per object and the less time spent
435 allocating/deallocating new links. The smaller the number, the less
436 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000437*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000438#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000439
440typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000441 PyObject_HEAD
442 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700443 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000444 PyObject *nextlink;
445 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000446} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000447
448typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000449 PyObject_HEAD
450 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700451 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000452 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000453} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000454
Raymond Hettingerad983e72003-11-12 14:32:26 +0000455static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000456
457static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000458teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000459{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000460 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000461
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000462 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
463 if (tdo == NULL)
464 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000466 tdo->numread = 0;
467 tdo->nextlink = NULL;
468 Py_INCREF(it);
469 tdo->it = it;
470 PyObject_GC_Track(tdo);
471 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472}
473
474static PyObject *
475teedataobject_jumplink(teedataobject *tdo)
476{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000477 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000478 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000479 Py_XINCREF(tdo->nextlink);
480 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000481}
482
483static PyObject *
484teedataobject_getitem(teedataobject *tdo, int i)
485{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000486 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000488 assert(i < LINKCELLS);
489 if (i < tdo->numread)
490 value = tdo->values[i];
491 else {
492 /* this is the lead iterator, so fetch more data */
493 assert(i == tdo->numread);
494 value = PyIter_Next(tdo->it);
495 if (value == NULL)
496 return NULL;
497 tdo->numread++;
498 tdo->values[i] = value;
499 }
500 Py_INCREF(value);
501 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000502}
503
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000504static int
505teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000507 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700508
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000509 Py_VISIT(tdo->it);
510 for (i = 0; i < tdo->numread; i++)
511 Py_VISIT(tdo->values[i]);
512 Py_VISIT(tdo->nextlink);
513 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000514}
515
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200516static void
517teedataobject_safe_decref(PyObject *obj)
518{
519 while (obj && Py_TYPE(obj) == &teedataobject_type &&
520 Py_REFCNT(obj) == 1) {
521 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
522 ((teedataobject *)obj)->nextlink = NULL;
523 Py_DECREF(obj);
524 obj = nextlink;
525 }
526 Py_XDECREF(obj);
527}
528
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000529static int
530teedataobject_clear(teedataobject *tdo)
531{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000532 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200533 PyObject *tmp;
534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000535 Py_CLEAR(tdo->it);
536 for (i=0 ; i<tdo->numread ; i++)
537 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200538 tmp = tdo->nextlink;
539 tdo->nextlink = NULL;
540 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000541 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000542}
543
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000544static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000545teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000546{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000547 PyObject_GC_UnTrack(tdo);
548 teedataobject_clear(tdo);
549 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550}
551
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000552static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530553teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000554{
555 int i;
556 /* create a temporary list of already iterated values */
557 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700558
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000559 if (!values)
560 return NULL;
561 for (i=0 ; i<tdo->numread ; i++) {
562 Py_INCREF(tdo->values[i]);
563 PyList_SET_ITEM(values, i, tdo->values[i]);
564 }
565 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
566 values,
567 tdo->nextlink ? tdo->nextlink : Py_None);
568}
569
Tal Einatc4bccd32018-09-12 00:49:13 +0300570/*[clinic input]
571@classmethod
572itertools.teedataobject.__new__
573 iterable as it: object
574 values: object(subclass_of='&PyList_Type')
575 next: object
576 /
577Data container common to multiple tee objects.
578[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000579
580static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300581itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
582 PyObject *values, PyObject *next)
583/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000584{
585 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000586 Py_ssize_t i, len;
587
588 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000589
590 tdo = (teedataobject *)teedataobject_newinternal(it);
591 if (!tdo)
592 return NULL;
593
594 len = PyList_GET_SIZE(values);
595 if (len > LINKCELLS)
596 goto err;
597 for (i=0; i<len; i++) {
598 tdo->values[i] = PyList_GET_ITEM(values, i);
599 Py_INCREF(tdo->values[i]);
600 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200601 /* len <= LINKCELLS < INT_MAX */
602 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000603
604 if (len == LINKCELLS) {
605 if (next != Py_None) {
606 if (Py_TYPE(next) != &teedataobject_type)
607 goto err;
608 assert(tdo->nextlink == NULL);
609 Py_INCREF(next);
610 tdo->nextlink = next;
611 }
612 } else {
613 if (next != Py_None)
614 goto err; /* shouldn't have a next if we are not full */
615 }
616 return (PyObject*)tdo;
617
618err:
619 Py_XDECREF(tdo);
620 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
621 return NULL;
622}
623
624static PyMethodDef teedataobject_methods[] = {
625 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
626 reduce_doc},
627 {NULL, NULL} /* sentinel */
628};
629
Raymond Hettingerad983e72003-11-12 14:32:26 +0000630static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700631 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000632 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 sizeof(teedataobject), /* tp_basicsize */
634 0, /* tp_itemsize */
635 /* methods */
636 (destructor)teedataobject_dealloc, /* tp_dealloc */
637 0, /* tp_print */
638 0, /* tp_getattr */
639 0, /* tp_setattr */
640 0, /* tp_reserved */
641 0, /* tp_repr */
642 0, /* tp_as_number */
643 0, /* tp_as_sequence */
644 0, /* tp_as_mapping */
645 0, /* tp_hash */
646 0, /* tp_call */
647 0, /* tp_str */
648 PyObject_GenericGetAttr, /* tp_getattro */
649 0, /* tp_setattro */
650 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700651 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300652 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000653 (traverseproc)teedataobject_traverse, /* tp_traverse */
654 (inquiry)teedataobject_clear, /* tp_clear */
655 0, /* tp_richcompare */
656 0, /* tp_weaklistoffset */
657 0, /* tp_iter */
658 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000659 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 0, /* tp_members */
661 0, /* tp_getset */
662 0, /* tp_base */
663 0, /* tp_dict */
664 0, /* tp_descr_get */
665 0, /* tp_descr_set */
666 0, /* tp_dictoffset */
667 0, /* tp_init */
668 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300669 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000670 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000671};
672
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000673
674static PyTypeObject tee_type;
675
676static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000677tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000678{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000680
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000681 if (to->index >= LINKCELLS) {
682 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200683 if (link == NULL)
684 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300685 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000686 to->index = 0;
687 }
688 value = teedataobject_getitem(to->dataobj, to->index);
689 if (value == NULL)
690 return NULL;
691 to->index++;
692 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000693}
694
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000695static int
696tee_traverse(teeobject *to, visitproc visit, void *arg)
697{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000698 Py_VISIT((PyObject *)to->dataobj);
699 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000700}
701
Raymond Hettingerad983e72003-11-12 14:32:26 +0000702static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530703tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000704{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000705 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000706
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000707 newto = PyObject_GC_New(teeobject, &tee_type);
708 if (newto == NULL)
709 return NULL;
710 Py_INCREF(to->dataobj);
711 newto->dataobj = to->dataobj;
712 newto->index = to->index;
713 newto->weakreflist = NULL;
714 PyObject_GC_Track(newto);
715 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000716}
717
718PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
719
720static PyObject *
721tee_fromiterable(PyObject *iterable)
722{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000723 teeobject *to;
724 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000725
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000726 it = PyObject_GetIter(iterable);
727 if (it == NULL)
728 return NULL;
729 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530730 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000731 goto done;
732 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000733
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000734 to = PyObject_GC_New(teeobject, &tee_type);
735 if (to == NULL)
736 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000737 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000738 if (!to->dataobj) {
739 PyObject_GC_Del(to);
740 to = NULL;
741 goto done;
742 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000743
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000744 to->index = 0;
745 to->weakreflist = NULL;
746 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000747done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000748 Py_XDECREF(it);
749 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000750}
751
Tal Einatc4bccd32018-09-12 00:49:13 +0300752/*[clinic input]
753@classmethod
754itertools._tee.__new__
755 iterable: object
756 /
757Iterator wrapped to make it copyable.
758[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000759
Tal Einatc4bccd32018-09-12 00:49:13 +0300760static PyObject *
761itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
762/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
763{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000764 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000765}
766
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000767static int
768tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000769{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000770 if (to->weakreflist != NULL)
771 PyObject_ClearWeakRefs((PyObject *) to);
772 Py_CLEAR(to->dataobj);
773 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000774}
775
776static void
777tee_dealloc(teeobject *to)
778{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000779 PyObject_GC_UnTrack(to);
780 tee_clear(to);
781 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000782}
783
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000784static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530785tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000786{
787 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
788}
789
790static PyObject *
791tee_setstate(teeobject *to, PyObject *state)
792{
793 teedataobject *tdo;
794 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300795 if (!PyTuple_Check(state)) {
796 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000797 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300798 }
799 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
800 return NULL;
801 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000802 if (index < 0 || index > LINKCELLS) {
803 PyErr_SetString(PyExc_ValueError, "Index out of range");
804 return NULL;
805 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200806 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300807 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000808 to->index = index;
809 Py_RETURN_NONE;
810}
811
Raymond Hettingerad983e72003-11-12 14:32:26 +0000812static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700813 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
814 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
815 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000816 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000817};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000818
819static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000820 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000821 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000822 sizeof(teeobject), /* tp_basicsize */
823 0, /* tp_itemsize */
824 /* methods */
825 (destructor)tee_dealloc, /* tp_dealloc */
826 0, /* tp_print */
827 0, /* tp_getattr */
828 0, /* tp_setattr */
829 0, /* tp_reserved */
830 0, /* tp_repr */
831 0, /* tp_as_number */
832 0, /* tp_as_sequence */
833 0, /* tp_as_mapping */
834 0, /* tp_hash */
835 0, /* tp_call */
836 0, /* tp_str */
837 0, /* tp_getattro */
838 0, /* tp_setattro */
839 0, /* tp_as_buffer */
840 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300841 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000842 (traverseproc)tee_traverse, /* tp_traverse */
843 (inquiry)tee_clear, /* tp_clear */
844 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700845 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000846 PyObject_SelfIter, /* tp_iter */
847 (iternextfunc)tee_next, /* tp_iternext */
848 tee_methods, /* tp_methods */
849 0, /* tp_members */
850 0, /* tp_getset */
851 0, /* tp_base */
852 0, /* tp_dict */
853 0, /* tp_descr_get */
854 0, /* tp_descr_set */
855 0, /* tp_dictoffset */
856 0, /* tp_init */
857 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300858 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000859 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000860};
861
Tal Einatc4bccd32018-09-12 00:49:13 +0300862/*[clinic input]
863itertools.tee
864 iterable: object
865 n: Py_ssize_t = 2
866 /
867Returns a tuple of n independent iterators.
868[clinic start generated code]*/
869
Raymond Hettingerad983e72003-11-12 14:32:26 +0000870static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300871itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
872/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000873{
Tal Einatc4bccd32018-09-12 00:49:13 +0300874 Py_ssize_t i;
875 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200876 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000878 if (n < 0) {
879 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
880 return NULL;
881 }
882 result = PyTuple_New(n);
883 if (result == NULL)
884 return NULL;
885 if (n == 0)
886 return result;
887 it = PyObject_GetIter(iterable);
888 if (it == NULL) {
889 Py_DECREF(result);
890 return NULL;
891 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200892
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200893 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200894 Py_DECREF(it);
895 Py_DECREF(result);
896 return NULL;
897 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200898 if (copyfunc != NULL) {
899 copyable = it;
900 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200901 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000902 copyable = tee_fromiterable(it);
903 Py_DECREF(it);
904 if (copyable == NULL) {
905 Py_DECREF(result);
906 return NULL;
907 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200908 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
909 if (copyfunc == NULL) {
910 Py_DECREF(copyable);
911 Py_DECREF(result);
912 return NULL;
913 }
914 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200915
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200916 PyTuple_SET_ITEM(result, 0, copyable);
917 for (i = 1; i < n; i++) {
918 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000919 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200920 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000921 Py_DECREF(result);
922 return NULL;
923 }
924 PyTuple_SET_ITEM(result, i, copyable);
925 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200926 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000927 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000928}
929
Raymond Hettingerad983e72003-11-12 14:32:26 +0000930
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700931/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000932
933typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 PyObject_HEAD
935 PyObject *it;
936 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700937 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000938 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000939} cycleobject;
940
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000941static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000942
Tal Einatc4bccd32018-09-12 00:49:13 +0300943/*[clinic input]
944@classmethod
945itertools.cycle.__new__
946 iterable: object
947 /
948Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
949[clinic start generated code]*/
950
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000951static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300952itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
953/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000954{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000955 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000956 PyObject *saved;
957 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000958
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000959 /* Get iterator. */
960 it = PyObject_GetIter(iterable);
961 if (it == NULL)
962 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000963
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000964 saved = PyList_New(0);
965 if (saved == NULL) {
966 Py_DECREF(it);
967 return NULL;
968 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000969
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000970 /* create cycleobject structure */
971 lz = (cycleobject *)type->tp_alloc(type, 0);
972 if (lz == NULL) {
973 Py_DECREF(it);
974 Py_DECREF(saved);
975 return NULL;
976 }
977 lz->it = it;
978 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700979 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000980 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000981
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000982 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000983}
984
985static void
986cycle_dealloc(cycleobject *lz)
987{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000988 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000989 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700990 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000991 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000992}
993
994static int
995cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
996{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700997 if (lz->it)
998 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000999 Py_VISIT(lz->saved);
1000 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001001}
1002
1003static PyObject *
1004cycle_next(cycleobject *lz)
1005{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001006 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001007
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001008 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001009 item = PyIter_Next(lz->it);
1010 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001011 if (lz->firstpass)
1012 return item;
1013 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001014 Py_DECREF(item);
1015 return NULL;
1016 }
1017 return item;
1018 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001019 /* Note: StopIteration is already cleared by PyIter_Next() */
1020 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001021 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001022 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001023 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001024 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001025 return NULL;
1026 item = PyList_GET_ITEM(lz->saved, lz->index);
1027 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001028 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001029 lz->index = 0;
1030 Py_INCREF(item);
1031 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001032}
1033
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001034static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301035cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001036{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001037 /* Create a new cycle with the iterator tuple, then set the saved state */
1038 if (lz->it == NULL) {
1039 PyObject *it = PyObject_GetIter(lz->saved);
1040 if (it == NULL)
1041 return NULL;
1042 if (lz->index != 0) {
1043 _Py_IDENTIFIER(__setstate__);
1044 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1045 "n", lz->index);
1046 if (res == NULL) {
1047 Py_DECREF(it);
1048 return NULL;
1049 }
1050 Py_DECREF(res);
1051 }
1052 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
1053 }
1054 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
1055 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001056}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001057
1058static PyObject *
1059cycle_setstate(cycleobject *lz, PyObject *state)
1060{
1061 PyObject *saved=NULL;
1062 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001063 if (!PyTuple_Check(state)) {
1064 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001065 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001066 }
1067 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1068 return NULL;
1069 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001070 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001071 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001072 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001073 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001074 Py_RETURN_NONE;
1075}
1076
1077static PyMethodDef cycle_methods[] = {
1078 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1079 reduce_doc},
1080 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1081 setstate_doc},
1082 {NULL, NULL} /* sentinel */
1083};
1084
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001085static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001086 PyVarObject_HEAD_INIT(NULL, 0)
1087 "itertools.cycle", /* tp_name */
1088 sizeof(cycleobject), /* tp_basicsize */
1089 0, /* tp_itemsize */
1090 /* methods */
1091 (destructor)cycle_dealloc, /* tp_dealloc */
1092 0, /* tp_print */
1093 0, /* tp_getattr */
1094 0, /* tp_setattr */
1095 0, /* tp_reserved */
1096 0, /* tp_repr */
1097 0, /* tp_as_number */
1098 0, /* tp_as_sequence */
1099 0, /* tp_as_mapping */
1100 0, /* tp_hash */
1101 0, /* tp_call */
1102 0, /* tp_str */
1103 PyObject_GenericGetAttr, /* tp_getattro */
1104 0, /* tp_setattro */
1105 0, /* tp_as_buffer */
1106 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1107 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001108 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001109 (traverseproc)cycle_traverse, /* tp_traverse */
1110 0, /* tp_clear */
1111 0, /* tp_richcompare */
1112 0, /* tp_weaklistoffset */
1113 PyObject_SelfIter, /* tp_iter */
1114 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001115 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001116 0, /* tp_members */
1117 0, /* tp_getset */
1118 0, /* tp_base */
1119 0, /* tp_dict */
1120 0, /* tp_descr_get */
1121 0, /* tp_descr_set */
1122 0, /* tp_dictoffset */
1123 0, /* tp_init */
1124 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001125 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001126 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001127};
1128
1129
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130/* dropwhile object **********************************************************/
1131
1132typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001133 PyObject_HEAD
1134 PyObject *func;
1135 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001136 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001137} dropwhileobject;
1138
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001139static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001140
Tal Einatc4bccd32018-09-12 00:49:13 +03001141/*[clinic input]
1142@classmethod
1143itertools.dropwhile.__new__
1144 predicate as func: object
1145 iterable as seq: object
1146 /
1147Drop items from the iterable while predicate(item) is true.
1148
1149Afterwards, return every element until the iterable is exhausted.
1150[clinic start generated code]*/
1151
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001152static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001153itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1154/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001155{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001156 PyObject *it;
1157 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001159 /* Get iterator. */
1160 it = PyObject_GetIter(seq);
1161 if (it == NULL)
1162 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001164 /* create dropwhileobject structure */
1165 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1166 if (lz == NULL) {
1167 Py_DECREF(it);
1168 return NULL;
1169 }
1170 Py_INCREF(func);
1171 lz->func = func;
1172 lz->it = it;
1173 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001174
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001175 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001176}
1177
1178static void
1179dropwhile_dealloc(dropwhileobject *lz)
1180{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001181 PyObject_GC_UnTrack(lz);
1182 Py_XDECREF(lz->func);
1183 Py_XDECREF(lz->it);
1184 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001185}
1186
1187static int
1188dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1189{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001190 Py_VISIT(lz->it);
1191 Py_VISIT(lz->func);
1192 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001193}
1194
1195static PyObject *
1196dropwhile_next(dropwhileobject *lz)
1197{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001198 PyObject *item, *good;
1199 PyObject *it = lz->it;
1200 long ok;
1201 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001203 iternext = *Py_TYPE(it)->tp_iternext;
1204 for (;;) {
1205 item = iternext(it);
1206 if (item == NULL)
1207 return NULL;
1208 if (lz->start == 1)
1209 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001211 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001212 if (good == NULL) {
1213 Py_DECREF(item);
1214 return NULL;
1215 }
1216 ok = PyObject_IsTrue(good);
1217 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001218 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001219 lz->start = 1;
1220 return item;
1221 }
1222 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001223 if (ok < 0)
1224 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001225 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001226}
1227
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001228static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301229dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001230{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001231 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 +00001232}
1233
1234static PyObject *
1235dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1236{
1237 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001238 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001239 return NULL;
1240 lz->start = start;
1241 Py_RETURN_NONE;
1242}
1243
1244static PyMethodDef dropwhile_methods[] = {
1245 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1246 reduce_doc},
1247 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1248 setstate_doc},
1249 {NULL, NULL} /* sentinel */
1250};
1251
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001252static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001253 PyVarObject_HEAD_INIT(NULL, 0)
1254 "itertools.dropwhile", /* tp_name */
1255 sizeof(dropwhileobject), /* tp_basicsize */
1256 0, /* tp_itemsize */
1257 /* methods */
1258 (destructor)dropwhile_dealloc, /* tp_dealloc */
1259 0, /* tp_print */
1260 0, /* tp_getattr */
1261 0, /* tp_setattr */
1262 0, /* tp_reserved */
1263 0, /* tp_repr */
1264 0, /* tp_as_number */
1265 0, /* tp_as_sequence */
1266 0, /* tp_as_mapping */
1267 0, /* tp_hash */
1268 0, /* tp_call */
1269 0, /* tp_str */
1270 PyObject_GenericGetAttr, /* tp_getattro */
1271 0, /* tp_setattro */
1272 0, /* tp_as_buffer */
1273 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1274 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001275 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001276 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001277 0, /* tp_clear */
1278 0, /* tp_richcompare */
1279 0, /* tp_weaklistoffset */
1280 PyObject_SelfIter, /* tp_iter */
1281 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001282 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001283 0, /* tp_members */
1284 0, /* tp_getset */
1285 0, /* tp_base */
1286 0, /* tp_dict */
1287 0, /* tp_descr_get */
1288 0, /* tp_descr_set */
1289 0, /* tp_dictoffset */
1290 0, /* tp_init */
1291 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001292 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001293 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294};
1295
1296
1297/* takewhile object **********************************************************/
1298
1299typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001300 PyObject_HEAD
1301 PyObject *func;
1302 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001303 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304} takewhileobject;
1305
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001306static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307
Tal Einatc4bccd32018-09-12 00:49:13 +03001308/*[clinic input]
1309@classmethod
1310itertools.takewhile.__new__
1311 predicate as func: object
1312 iterable as seq: object
1313 /
1314Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1315[clinic start generated code]*/
1316
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001317static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001318itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1319/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001320{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001321 PyObject *it;
1322 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001323
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001324 /* Get iterator. */
1325 it = PyObject_GetIter(seq);
1326 if (it == NULL)
1327 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001329 /* create takewhileobject structure */
1330 lz = (takewhileobject *)type->tp_alloc(type, 0);
1331 if (lz == NULL) {
1332 Py_DECREF(it);
1333 return NULL;
1334 }
1335 Py_INCREF(func);
1336 lz->func = func;
1337 lz->it = it;
1338 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001340 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341}
1342
1343static void
1344takewhile_dealloc(takewhileobject *lz)
1345{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001346 PyObject_GC_UnTrack(lz);
1347 Py_XDECREF(lz->func);
1348 Py_XDECREF(lz->it);
1349 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350}
1351
1352static int
1353takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1354{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001355 Py_VISIT(lz->it);
1356 Py_VISIT(lz->func);
1357 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001358}
1359
1360static PyObject *
1361takewhile_next(takewhileobject *lz)
1362{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 PyObject *item, *good;
1364 PyObject *it = lz->it;
1365 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001367 if (lz->stop == 1)
1368 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001369
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001370 item = (*Py_TYPE(it)->tp_iternext)(it);
1371 if (item == NULL)
1372 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001373
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001374 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001375 if (good == NULL) {
1376 Py_DECREF(item);
1377 return NULL;
1378 }
1379 ok = PyObject_IsTrue(good);
1380 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001381 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001382 return item;
1383 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001384 if (ok == 0)
1385 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001386 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387}
1388
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001389static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301390takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001391{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001392 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 +00001393}
1394
1395static PyObject *
1396takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1397{
1398 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001399
Antoine Pitrou721738f2012-08-15 23:20:39 +02001400 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001401 return NULL;
1402 lz->stop = stop;
1403 Py_RETURN_NONE;
1404}
1405
1406static PyMethodDef takewhile_reduce_methods[] = {
1407 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1408 reduce_doc},
1409 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1410 setstate_doc},
1411 {NULL, NULL} /* sentinel */
1412};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001413
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001414static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001415 PyVarObject_HEAD_INIT(NULL, 0)
1416 "itertools.takewhile", /* tp_name */
1417 sizeof(takewhileobject), /* tp_basicsize */
1418 0, /* tp_itemsize */
1419 /* methods */
1420 (destructor)takewhile_dealloc, /* tp_dealloc */
1421 0, /* tp_print */
1422 0, /* tp_getattr */
1423 0, /* tp_setattr */
1424 0, /* tp_reserved */
1425 0, /* tp_repr */
1426 0, /* tp_as_number */
1427 0, /* tp_as_sequence */
1428 0, /* tp_as_mapping */
1429 0, /* tp_hash */
1430 0, /* tp_call */
1431 0, /* tp_str */
1432 PyObject_GenericGetAttr, /* tp_getattro */
1433 0, /* tp_setattro */
1434 0, /* tp_as_buffer */
1435 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1436 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001437 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001438 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001439 0, /* tp_clear */
1440 0, /* tp_richcompare */
1441 0, /* tp_weaklistoffset */
1442 PyObject_SelfIter, /* tp_iter */
1443 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001444 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001445 0, /* tp_members */
1446 0, /* tp_getset */
1447 0, /* tp_base */
1448 0, /* tp_dict */
1449 0, /* tp_descr_get */
1450 0, /* tp_descr_set */
1451 0, /* tp_dictoffset */
1452 0, /* tp_init */
1453 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001454 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001455 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456};
1457
1458
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001459/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460
1461typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001462 PyObject_HEAD
1463 PyObject *it;
1464 Py_ssize_t next;
1465 Py_ssize_t stop;
1466 Py_ssize_t step;
1467 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468} isliceobject;
1469
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001470static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471
1472static PyObject *
1473islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1474{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001475 PyObject *seq;
1476 Py_ssize_t start=0, stop=-1, step=1;
1477 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1478 Py_ssize_t numargs;
1479 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001481 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001483
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001484 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1485 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001487 numargs = PyTuple_Size(args);
1488 if (numargs == 2) {
1489 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001490 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001491 if (stop == -1) {
1492 if (PyErr_Occurred())
1493 PyErr_Clear();
1494 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001495 "Stop argument for islice() must be None or "
1496 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001497 return NULL;
1498 }
1499 }
1500 } else {
1501 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001502 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001503 if (start == -1 && PyErr_Occurred())
1504 PyErr_Clear();
1505 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001506 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001507 if (stop == -1) {
1508 if (PyErr_Occurred())
1509 PyErr_Clear();
1510 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001511 "Stop argument for islice() must be None or "
1512 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001513 return NULL;
1514 }
1515 }
1516 }
1517 if (start<0 || stop<-1) {
1518 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001519 "Indices for islice() must be None or "
1520 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001521 return NULL;
1522 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001524 if (a3 != NULL) {
1525 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001526 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001527 if (step == -1 && PyErr_Occurred())
1528 PyErr_Clear();
1529 }
1530 if (step<1) {
1531 PyErr_SetString(PyExc_ValueError,
1532 "Step for islice() must be a positive integer or None.");
1533 return NULL;
1534 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001536 /* Get iterator. */
1537 it = PyObject_GetIter(seq);
1538 if (it == NULL)
1539 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001541 /* create isliceobject structure */
1542 lz = (isliceobject *)type->tp_alloc(type, 0);
1543 if (lz == NULL) {
1544 Py_DECREF(it);
1545 return NULL;
1546 }
1547 lz->it = it;
1548 lz->next = start;
1549 lz->stop = stop;
1550 lz->step = step;
1551 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001552
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001553 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554}
1555
1556static void
1557islice_dealloc(isliceobject *lz)
1558{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001559 PyObject_GC_UnTrack(lz);
1560 Py_XDECREF(lz->it);
1561 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562}
1563
1564static int
1565islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1566{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001567 Py_VISIT(lz->it);
1568 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001569}
1570
1571static PyObject *
1572islice_next(isliceobject *lz)
1573{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001574 PyObject *item;
1575 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001576 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001577 Py_ssize_t oldnext;
1578 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001579
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001580 if (it == NULL)
1581 return NULL;
1582
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001583 iternext = *Py_TYPE(it)->tp_iternext;
1584 while (lz->cnt < lz->next) {
1585 item = iternext(it);
1586 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001587 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001588 Py_DECREF(item);
1589 lz->cnt++;
1590 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001591 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001592 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001593 item = iternext(it);
1594 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001595 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001596 lz->cnt++;
1597 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001598 /* The (size_t) cast below avoids the danger of undefined
1599 behaviour from signed integer overflow. */
1600 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001601 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1602 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001603 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001604
1605empty:
1606 Py_CLEAR(lz->it);
1607 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608}
1609
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001610static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301611islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001612{
1613 /* When unpickled, generate a new object with the same bounds,
1614 * then 'setstate' with the next and count
1615 */
1616 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001617
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001618 if (lz->it == NULL) {
1619 PyObject *empty_list;
1620 PyObject *empty_it;
1621 empty_list = PyList_New(0);
1622 if (empty_list == NULL)
1623 return NULL;
1624 empty_it = PyObject_GetIter(empty_list);
1625 Py_DECREF(empty_list);
1626 if (empty_it == NULL)
1627 return NULL;
1628 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1629 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001630 if (lz->stop == -1) {
1631 stop = Py_None;
1632 Py_INCREF(stop);
1633 } else {
1634 stop = PyLong_FromSsize_t(lz->stop);
1635 if (stop == NULL)
1636 return NULL;
1637 }
1638 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1639 lz->it, lz->next, stop, lz->step,
1640 lz->cnt);
1641}
1642
1643static PyObject *
1644islice_setstate(isliceobject *lz, PyObject *state)
1645{
1646 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001647
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001648 if (cnt == -1 && PyErr_Occurred())
1649 return NULL;
1650 lz->cnt = cnt;
1651 Py_RETURN_NONE;
1652}
1653
1654static PyMethodDef islice_methods[] = {
1655 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1656 reduce_doc},
1657 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1658 setstate_doc},
1659 {NULL, NULL} /* sentinel */
1660};
1661
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001663"islice(iterable, stop) --> islice object\n\
1664islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665\n\
1666Return an iterator whose next() method returns selected values from an\n\
1667iterable. If start is specified, will skip all preceding elements;\n\
1668otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001669specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670skipped between successive calls. Works like a slice() on a list\n\
1671but returns an iterator.");
1672
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001673static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001674 PyVarObject_HEAD_INIT(NULL, 0)
1675 "itertools.islice", /* tp_name */
1676 sizeof(isliceobject), /* tp_basicsize */
1677 0, /* tp_itemsize */
1678 /* methods */
1679 (destructor)islice_dealloc, /* tp_dealloc */
1680 0, /* tp_print */
1681 0, /* tp_getattr */
1682 0, /* tp_setattr */
1683 0, /* tp_reserved */
1684 0, /* tp_repr */
1685 0, /* tp_as_number */
1686 0, /* tp_as_sequence */
1687 0, /* tp_as_mapping */
1688 0, /* tp_hash */
1689 0, /* tp_call */
1690 0, /* tp_str */
1691 PyObject_GenericGetAttr, /* tp_getattro */
1692 0, /* tp_setattro */
1693 0, /* tp_as_buffer */
1694 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1695 Py_TPFLAGS_BASETYPE, /* tp_flags */
1696 islice_doc, /* tp_doc */
1697 (traverseproc)islice_traverse, /* tp_traverse */
1698 0, /* tp_clear */
1699 0, /* tp_richcompare */
1700 0, /* tp_weaklistoffset */
1701 PyObject_SelfIter, /* tp_iter */
1702 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001703 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001704 0, /* tp_members */
1705 0, /* tp_getset */
1706 0, /* tp_base */
1707 0, /* tp_dict */
1708 0, /* tp_descr_get */
1709 0, /* tp_descr_set */
1710 0, /* tp_dictoffset */
1711 0, /* tp_init */
1712 0, /* tp_alloc */
1713 islice_new, /* tp_new */
1714 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001715};
1716
1717
1718/* starmap object ************************************************************/
1719
1720typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001721 PyObject_HEAD
1722 PyObject *func;
1723 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724} starmapobject;
1725
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001726static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727
Tal Einatc4bccd32018-09-12 00:49:13 +03001728/*[clinic input]
1729@classmethod
1730itertools.starmap.__new__
1731 function as func: object
1732 iterable as seq: object
1733 /
1734Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1735[clinic start generated code]*/
1736
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001738itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1739/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001740{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001741 PyObject *it;
1742 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001743
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001744 /* Get iterator. */
1745 it = PyObject_GetIter(seq);
1746 if (it == NULL)
1747 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001748
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001749 /* create starmapobject structure */
1750 lz = (starmapobject *)type->tp_alloc(type, 0);
1751 if (lz == NULL) {
1752 Py_DECREF(it);
1753 return NULL;
1754 }
1755 Py_INCREF(func);
1756 lz->func = func;
1757 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001759 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001760}
1761
1762static void
1763starmap_dealloc(starmapobject *lz)
1764{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001765 PyObject_GC_UnTrack(lz);
1766 Py_XDECREF(lz->func);
1767 Py_XDECREF(lz->it);
1768 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001769}
1770
1771static int
1772starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1773{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001774 Py_VISIT(lz->it);
1775 Py_VISIT(lz->func);
1776 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001777}
1778
1779static PyObject *
1780starmap_next(starmapobject *lz)
1781{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001782 PyObject *args;
1783 PyObject *result;
1784 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001785
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001786 args = (*Py_TYPE(it)->tp_iternext)(it);
1787 if (args == NULL)
1788 return NULL;
1789 if (!PyTuple_CheckExact(args)) {
1790 PyObject *newargs = PySequence_Tuple(args);
1791 Py_DECREF(args);
1792 if (newargs == NULL)
1793 return NULL;
1794 args = newargs;
1795 }
1796 result = PyObject_Call(lz->func, args, NULL);
1797 Py_DECREF(args);
1798 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001799}
1800
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001801static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301802starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001803{
1804 /* Just pickle the iterator */
1805 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1806}
1807
1808static PyMethodDef starmap_methods[] = {
1809 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1810 reduce_doc},
1811 {NULL, NULL} /* sentinel */
1812};
1813
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001814static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001815 PyVarObject_HEAD_INIT(NULL, 0)
1816 "itertools.starmap", /* tp_name */
1817 sizeof(starmapobject), /* tp_basicsize */
1818 0, /* tp_itemsize */
1819 /* methods */
1820 (destructor)starmap_dealloc, /* tp_dealloc */
1821 0, /* tp_print */
1822 0, /* tp_getattr */
1823 0, /* tp_setattr */
1824 0, /* tp_reserved */
1825 0, /* tp_repr */
1826 0, /* tp_as_number */
1827 0, /* tp_as_sequence */
1828 0, /* tp_as_mapping */
1829 0, /* tp_hash */
1830 0, /* tp_call */
1831 0, /* tp_str */
1832 PyObject_GenericGetAttr, /* tp_getattro */
1833 0, /* tp_setattro */
1834 0, /* tp_as_buffer */
1835 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1836 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001837 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001838 (traverseproc)starmap_traverse, /* tp_traverse */
1839 0, /* tp_clear */
1840 0, /* tp_richcompare */
1841 0, /* tp_weaklistoffset */
1842 PyObject_SelfIter, /* tp_iter */
1843 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001844 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001845 0, /* tp_members */
1846 0, /* tp_getset */
1847 0, /* tp_base */
1848 0, /* tp_dict */
1849 0, /* tp_descr_get */
1850 0, /* tp_descr_set */
1851 0, /* tp_dictoffset */
1852 0, /* tp_init */
1853 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001854 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001855 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001856};
1857
1858
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001859/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001860
1861typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001862 PyObject_HEAD
1863 PyObject *source; /* Iterator over input iterables */
1864 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001865} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001866
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001867static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001868
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001869static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001870chain_new_internal(PyTypeObject *type, PyObject *source)
1871{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001872 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001873
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001874 lz = (chainobject *)type->tp_alloc(type, 0);
1875 if (lz == NULL) {
1876 Py_DECREF(source);
1877 return NULL;
1878 }
1879
1880 lz->source = source;
1881 lz->active = NULL;
1882 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001883}
1884
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001885static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001886chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001887{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001888 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001889
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001890 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001891 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001892
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001893 source = PyObject_GetIter(args);
1894 if (source == NULL)
1895 return NULL;
1896
1897 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001898}
1899
Tal Einatc4bccd32018-09-12 00:49:13 +03001900/*[clinic input]
1901@classmethod
1902itertools.chain.from_iterable
1903 iterable as arg: object
1904 /
1905Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1906[clinic start generated code]*/
1907
Christian Heimesf16baeb2008-02-29 14:57:44 +00001908static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001909itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1910/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001911{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001912 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001913
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001914 source = PyObject_GetIter(arg);
1915 if (source == NULL)
1916 return NULL;
1917
1918 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001919}
1920
1921static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001922chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001923{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001924 PyObject_GC_UnTrack(lz);
1925 Py_XDECREF(lz->active);
1926 Py_XDECREF(lz->source);
1927 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001928}
1929
Raymond Hettinger2012f172003-02-07 05:32:58 +00001930static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001931chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001932{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001933 Py_VISIT(lz->source);
1934 Py_VISIT(lz->active);
1935 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001936}
1937
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001938static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001939chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001940{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001941 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001942
T. Wouters5466d4a2017-03-30 09:58:35 -07001943 /* lz->source is the iterator of iterables. If it's NULL, we've already
1944 * consumed them all. lz->active is the current iterator. If it's NULL,
1945 * we should grab a new one from lz->source. */
1946 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001947 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001948 PyObject *iterable = PyIter_Next(lz->source);
1949 if (iterable == NULL) {
1950 Py_CLEAR(lz->source);
1951 return NULL; /* no more input sources */
1952 }
1953 lz->active = PyObject_GetIter(iterable);
1954 Py_DECREF(iterable);
1955 if (lz->active == NULL) {
1956 Py_CLEAR(lz->source);
1957 return NULL; /* input not iterable */
1958 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001959 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001960 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1961 if (item != NULL)
1962 return item;
1963 if (PyErr_Occurred()) {
1964 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1965 PyErr_Clear();
1966 else
1967 return NULL; /* input raised an exception */
1968 }
1969 /* lz->active is consumed, try with the next iterable. */
1970 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001971 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001972 /* Everything had been consumed already. */
1973 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001974}
1975
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001976static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301977chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001978{
1979 if (lz->source) {
1980 /* we can't pickle function objects (itertools.from_iterable) so
1981 * we must use setstate to replace the iterable. One day we
1982 * will fix pickling of functions
1983 */
1984 if (lz->active) {
1985 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1986 } else {
1987 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1988 }
1989 } else {
1990 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1991 }
1992 return NULL;
1993}
1994
1995static PyObject *
1996chain_setstate(chainobject *lz, PyObject *state)
1997{
1998 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001999
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002000 if (!PyTuple_Check(state)) {
2001 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002002 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002003 }
2004 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2005 return NULL;
2006 }
2007 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2008 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2009 return NULL;
2010 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002011
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002012 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002013 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002014 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002015 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002016 Py_RETURN_NONE;
2017}
2018
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002019PyDoc_STRVAR(chain_doc,
2020"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002021\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002022Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002023first iterable until it is exhausted, then elements from the next\n\
2024iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002025
Christian Heimesf16baeb2008-02-29 14:57:44 +00002026static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002027 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002028 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2029 reduce_doc},
2030 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2031 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002032 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002033};
2034
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002035static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002036 PyVarObject_HEAD_INIT(NULL, 0)
2037 "itertools.chain", /* tp_name */
2038 sizeof(chainobject), /* tp_basicsize */
2039 0, /* tp_itemsize */
2040 /* methods */
2041 (destructor)chain_dealloc, /* tp_dealloc */
2042 0, /* tp_print */
2043 0, /* tp_getattr */
2044 0, /* tp_setattr */
2045 0, /* tp_reserved */
2046 0, /* tp_repr */
2047 0, /* tp_as_number */
2048 0, /* tp_as_sequence */
2049 0, /* tp_as_mapping */
2050 0, /* tp_hash */
2051 0, /* tp_call */
2052 0, /* tp_str */
2053 PyObject_GenericGetAttr, /* tp_getattro */
2054 0, /* tp_setattro */
2055 0, /* tp_as_buffer */
2056 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2057 Py_TPFLAGS_BASETYPE, /* tp_flags */
2058 chain_doc, /* tp_doc */
2059 (traverseproc)chain_traverse, /* tp_traverse */
2060 0, /* tp_clear */
2061 0, /* tp_richcompare */
2062 0, /* tp_weaklistoffset */
2063 PyObject_SelfIter, /* tp_iter */
2064 (iternextfunc)chain_next, /* tp_iternext */
2065 chain_methods, /* tp_methods */
2066 0, /* tp_members */
2067 0, /* tp_getset */
2068 0, /* tp_base */
2069 0, /* tp_dict */
2070 0, /* tp_descr_get */
2071 0, /* tp_descr_set */
2072 0, /* tp_dictoffset */
2073 0, /* tp_init */
2074 0, /* tp_alloc */
2075 chain_new, /* tp_new */
2076 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002077};
2078
2079
Christian Heimesc3f30c42008-02-22 16:37:40 +00002080/* product object ************************************************************/
2081
2082typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002083 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002084 PyObject *pools; /* tuple of pool tuples */
2085 Py_ssize_t *indices; /* one index per pool */
2086 PyObject *result; /* most recently returned result tuple */
2087 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002088} productobject;
2089
2090static PyTypeObject product_type;
2091
2092static PyObject *
2093product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2094{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002095 productobject *lz;
2096 Py_ssize_t nargs, npools, repeat=1;
2097 PyObject *pools = NULL;
2098 Py_ssize_t *indices = NULL;
2099 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002100
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002101 if (kwds != NULL) {
2102 char *kwlist[] = {"repeat", 0};
2103 PyObject *tmpargs = PyTuple_New(0);
2104 if (tmpargs == NULL)
2105 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002106 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2107 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002108 Py_DECREF(tmpargs);
2109 return NULL;
2110 }
2111 Py_DECREF(tmpargs);
2112 if (repeat < 0) {
2113 PyErr_SetString(PyExc_ValueError,
2114 "repeat argument cannot be negative");
2115 return NULL;
2116 }
2117 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002118
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002119 assert(PyTuple_CheckExact(args));
2120 if (repeat == 0) {
2121 nargs = 0;
2122 } else {
2123 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002124 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002125 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2126 return NULL;
2127 }
2128 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002129 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002130
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002131 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002132 if (indices == NULL) {
2133 PyErr_NoMemory();
2134 goto error;
2135 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002136
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002137 pools = PyTuple_New(npools);
2138 if (pools == NULL)
2139 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002141 for (i=0; i < nargs ; ++i) {
2142 PyObject *item = PyTuple_GET_ITEM(args, i);
2143 PyObject *pool = PySequence_Tuple(item);
2144 if (pool == NULL)
2145 goto error;
2146 PyTuple_SET_ITEM(pools, i, pool);
2147 indices[i] = 0;
2148 }
2149 for ( ; i < npools; ++i) {
2150 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2151 Py_INCREF(pool);
2152 PyTuple_SET_ITEM(pools, i, pool);
2153 indices[i] = 0;
2154 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002155
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002156 /* create productobject structure */
2157 lz = (productobject *)type->tp_alloc(type, 0);
2158 if (lz == NULL)
2159 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002160
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002161 lz->pools = pools;
2162 lz->indices = indices;
2163 lz->result = NULL;
2164 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002166 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002167
2168error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002169 if (indices != NULL)
2170 PyMem_Free(indices);
2171 Py_XDECREF(pools);
2172 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002173}
2174
2175static void
2176product_dealloc(productobject *lz)
2177{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002178 PyObject_GC_UnTrack(lz);
2179 Py_XDECREF(lz->pools);
2180 Py_XDECREF(lz->result);
2181 if (lz->indices != NULL)
2182 PyMem_Free(lz->indices);
2183 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002184}
2185
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002186static PyObject *
2187product_sizeof(productobject *lz, void *unused)
2188{
2189 Py_ssize_t res;
2190
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002191 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002192 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2193 return PyLong_FromSsize_t(res);
2194}
2195
2196PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2197
Christian Heimesc3f30c42008-02-22 16:37:40 +00002198static int
2199product_traverse(productobject *lz, visitproc visit, void *arg)
2200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002201 Py_VISIT(lz->pools);
2202 Py_VISIT(lz->result);
2203 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002204}
2205
2206static PyObject *
2207product_next(productobject *lz)
2208{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002209 PyObject *pool;
2210 PyObject *elem;
2211 PyObject *oldelem;
2212 PyObject *pools = lz->pools;
2213 PyObject *result = lz->result;
2214 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2215 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002216
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002217 if (lz->stopped)
2218 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002219
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002220 if (result == NULL) {
2221 /* On the first pass, return an initial tuple filled with the
2222 first element from each pool. */
2223 result = PyTuple_New(npools);
2224 if (result == NULL)
2225 goto empty;
2226 lz->result = result;
2227 for (i=0; i < npools; i++) {
2228 pool = PyTuple_GET_ITEM(pools, i);
2229 if (PyTuple_GET_SIZE(pool) == 0)
2230 goto empty;
2231 elem = PyTuple_GET_ITEM(pool, 0);
2232 Py_INCREF(elem);
2233 PyTuple_SET_ITEM(result, i, elem);
2234 }
2235 } else {
2236 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002237
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002238 /* Copy the previous result tuple or re-use it if available */
2239 if (Py_REFCNT(result) > 1) {
2240 PyObject *old_result = result;
2241 result = PyTuple_New(npools);
2242 if (result == NULL)
2243 goto empty;
2244 lz->result = result;
2245 for (i=0; i < npools; i++) {
2246 elem = PyTuple_GET_ITEM(old_result, i);
2247 Py_INCREF(elem);
2248 PyTuple_SET_ITEM(result, i, elem);
2249 }
2250 Py_DECREF(old_result);
2251 }
2252 /* Now, we've got the only copy so we can update it in-place */
2253 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002254
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002255 /* Update the pool indices right-to-left. Only advance to the
2256 next pool when the previous one rolls-over */
2257 for (i=npools-1 ; i >= 0 ; i--) {
2258 pool = PyTuple_GET_ITEM(pools, i);
2259 indices[i]++;
2260 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2261 /* Roll-over and advance to next pool */
2262 indices[i] = 0;
2263 elem = PyTuple_GET_ITEM(pool, 0);
2264 Py_INCREF(elem);
2265 oldelem = PyTuple_GET_ITEM(result, i);
2266 PyTuple_SET_ITEM(result, i, elem);
2267 Py_DECREF(oldelem);
2268 } else {
2269 /* No rollover. Just increment and stop here. */
2270 elem = PyTuple_GET_ITEM(pool, indices[i]);
2271 Py_INCREF(elem);
2272 oldelem = PyTuple_GET_ITEM(result, i);
2273 PyTuple_SET_ITEM(result, i, elem);
2274 Py_DECREF(oldelem);
2275 break;
2276 }
2277 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002278
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002279 /* If i is negative, then the indices have all rolled-over
2280 and we're done. */
2281 if (i < 0)
2282 goto empty;
2283 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002284
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002285 Py_INCREF(result);
2286 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002287
2288empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002289 lz->stopped = 1;
2290 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002291}
2292
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002293static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302294product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002295{
2296 if (lz->stopped) {
2297 return Py_BuildValue("O(())", Py_TYPE(lz));
2298 } else if (lz->result == NULL) {
2299 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2300 } else {
2301 PyObject *indices;
2302 Py_ssize_t n, i;
2303
2304 /* we must pickle the indices use them for setstate, and
2305 * additionally indicate that the iterator has started
2306 */
2307 n = PyTuple_GET_SIZE(lz->pools);
2308 indices = PyTuple_New(n);
2309 if (indices == NULL)
2310 return NULL;
2311 for (i=0; i<n; i++){
2312 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2313 if (!index) {
2314 Py_DECREF(indices);
2315 return NULL;
2316 }
2317 PyTuple_SET_ITEM(indices, i, index);
2318 }
2319 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2320 }
2321}
2322
2323static PyObject *
2324product_setstate(productobject *lz, PyObject *state)
2325{
2326 PyObject *result;
2327 Py_ssize_t n, i;
2328
2329 n = PyTuple_GET_SIZE(lz->pools);
2330 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2331 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2332 return NULL;
2333 }
2334 for (i=0; i<n; i++)
2335 {
2336 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2337 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002338 PyObject* pool;
2339 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002340 if (index < 0 && PyErr_Occurred())
2341 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002342 pool = PyTuple_GET_ITEM(lz->pools, i);
2343 poolsize = PyTuple_GET_SIZE(pool);
2344 if (poolsize == 0) {
2345 lz->stopped = 1;
2346 Py_RETURN_NONE;
2347 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002348 /* clamp the index */
2349 if (index < 0)
2350 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002351 else if (index > poolsize-1)
2352 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002353 lz->indices[i] = index;
2354 }
2355
2356 result = PyTuple_New(n);
2357 if (!result)
2358 return NULL;
2359 for (i=0; i<n; i++) {
2360 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2361 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2362 Py_INCREF(element);
2363 PyTuple_SET_ITEM(result, i, element);
2364 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002365 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002366 Py_RETURN_NONE;
2367}
2368
2369static PyMethodDef product_methods[] = {
2370 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2371 reduce_doc},
2372 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2373 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002374 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2375 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002376 {NULL, NULL} /* sentinel */
2377};
2378
Christian Heimesc3f30c42008-02-22 16:37:40 +00002379PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002380"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002381\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002382Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002383For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2384The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2385cycle in a manner similar to an odometer (with the rightmost element changing\n\
2386on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002387To compute the product of an iterable with itself, specify the number\n\
2388of repetitions with the optional repeat keyword argument. For example,\n\
2389product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002390product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2391product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2392
2393static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002394 PyVarObject_HEAD_INIT(NULL, 0)
2395 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002396 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002397 0, /* tp_itemsize */
2398 /* methods */
2399 (destructor)product_dealloc, /* tp_dealloc */
2400 0, /* tp_print */
2401 0, /* tp_getattr */
2402 0, /* tp_setattr */
2403 0, /* tp_reserved */
2404 0, /* tp_repr */
2405 0, /* tp_as_number */
2406 0, /* tp_as_sequence */
2407 0, /* tp_as_mapping */
2408 0, /* tp_hash */
2409 0, /* tp_call */
2410 0, /* tp_str */
2411 PyObject_GenericGetAttr, /* tp_getattro */
2412 0, /* tp_setattro */
2413 0, /* tp_as_buffer */
2414 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2415 Py_TPFLAGS_BASETYPE, /* tp_flags */
2416 product_doc, /* tp_doc */
2417 (traverseproc)product_traverse, /* tp_traverse */
2418 0, /* tp_clear */
2419 0, /* tp_richcompare */
2420 0, /* tp_weaklistoffset */
2421 PyObject_SelfIter, /* tp_iter */
2422 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002423 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002424 0, /* tp_members */
2425 0, /* tp_getset */
2426 0, /* tp_base */
2427 0, /* tp_dict */
2428 0, /* tp_descr_get */
2429 0, /* tp_descr_set */
2430 0, /* tp_dictoffset */
2431 0, /* tp_init */
2432 0, /* tp_alloc */
2433 product_new, /* tp_new */
2434 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002435};
2436
2437
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002438/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002439
2440typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002441 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002442 PyObject *pool; /* input converted to a tuple */
2443 Py_ssize_t *indices; /* one index per result element */
2444 PyObject *result; /* most recently returned result tuple */
2445 Py_ssize_t r; /* size of result tuple */
2446 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002447} combinationsobject;
2448
2449static PyTypeObject combinations_type;
2450
Tal Einatc4bccd32018-09-12 00:49:13 +03002451
2452/*[clinic input]
2453@classmethod
2454itertools.combinations.__new__
2455 iterable: object
2456 r: Py_ssize_t
2457Return successive r-length combinations of elements in the iterable.
2458
2459combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2460[clinic start generated code]*/
2461
Christian Heimes380f7f22008-02-28 11:19:05 +00002462static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002463itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2464 Py_ssize_t r)
2465/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002466{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002467 combinationsobject *co;
2468 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002469 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002470 Py_ssize_t *indices = NULL;
2471 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002472
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002473 pool = PySequence_Tuple(iterable);
2474 if (pool == NULL)
2475 goto error;
2476 n = PyTuple_GET_SIZE(pool);
2477 if (r < 0) {
2478 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2479 goto error;
2480 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002481
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002482 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002483 if (indices == NULL) {
2484 PyErr_NoMemory();
2485 goto error;
2486 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002488 for (i=0 ; i<r ; i++)
2489 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002490
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002491 /* create combinationsobject structure */
2492 co = (combinationsobject *)type->tp_alloc(type, 0);
2493 if (co == NULL)
2494 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002495
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002496 co->pool = pool;
2497 co->indices = indices;
2498 co->result = NULL;
2499 co->r = r;
2500 co->stopped = r > n ? 1 : 0;
2501
2502 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002503
2504error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002505 if (indices != NULL)
2506 PyMem_Free(indices);
2507 Py_XDECREF(pool);
2508 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002509}
2510
2511static void
2512combinations_dealloc(combinationsobject *co)
2513{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002514 PyObject_GC_UnTrack(co);
2515 Py_XDECREF(co->pool);
2516 Py_XDECREF(co->result);
2517 if (co->indices != NULL)
2518 PyMem_Free(co->indices);
2519 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002520}
2521
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002522static PyObject *
2523combinations_sizeof(combinationsobject *co, void *unused)
2524{
2525 Py_ssize_t res;
2526
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002527 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002528 res += co->r * sizeof(Py_ssize_t);
2529 return PyLong_FromSsize_t(res);
2530}
2531
Christian Heimes380f7f22008-02-28 11:19:05 +00002532static int
2533combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2534{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002535 Py_VISIT(co->pool);
2536 Py_VISIT(co->result);
2537 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002538}
2539
2540static PyObject *
2541combinations_next(combinationsobject *co)
2542{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002543 PyObject *elem;
2544 PyObject *oldelem;
2545 PyObject *pool = co->pool;
2546 Py_ssize_t *indices = co->indices;
2547 PyObject *result = co->result;
2548 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2549 Py_ssize_t r = co->r;
2550 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002551
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002552 if (co->stopped)
2553 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002554
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002555 if (result == NULL) {
2556 /* On the first pass, initialize result tuple using the indices */
2557 result = PyTuple_New(r);
2558 if (result == NULL)
2559 goto empty;
2560 co->result = result;
2561 for (i=0; i<r ; i++) {
2562 index = indices[i];
2563 elem = PyTuple_GET_ITEM(pool, index);
2564 Py_INCREF(elem);
2565 PyTuple_SET_ITEM(result, i, elem);
2566 }
2567 } else {
2568 /* Copy the previous result tuple or re-use it if available */
2569 if (Py_REFCNT(result) > 1) {
2570 PyObject *old_result = result;
2571 result = PyTuple_New(r);
2572 if (result == NULL)
2573 goto empty;
2574 co->result = result;
2575 for (i=0; i<r ; i++) {
2576 elem = PyTuple_GET_ITEM(old_result, i);
2577 Py_INCREF(elem);
2578 PyTuple_SET_ITEM(result, i, elem);
2579 }
2580 Py_DECREF(old_result);
2581 }
2582 /* Now, we've got the only copy so we can update it in-place
2583 * CPython's empty tuple is a singleton and cached in
2584 * PyTuple's freelist.
2585 */
2586 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002587
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002588 /* Scan indices right-to-left until finding one that is not
2589 at its maximum (i + n - r). */
2590 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2591 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002592
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002593 /* If i is negative, then the indices are all at
2594 their maximum value and we're done. */
2595 if (i < 0)
2596 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002598 /* Increment the current index which we know is not at its
2599 maximum. Then move back to the right setting each index
2600 to its lowest possible value (one higher than the index
2601 to its left -- this maintains the sort order invariant). */
2602 indices[i]++;
2603 for (j=i+1 ; j<r ; j++)
2604 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002605
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002606 /* Update the result tuple for the new indices
2607 starting with i, the leftmost index that changed */
2608 for ( ; i<r ; i++) {
2609 index = indices[i];
2610 elem = PyTuple_GET_ITEM(pool, index);
2611 Py_INCREF(elem);
2612 oldelem = PyTuple_GET_ITEM(result, i);
2613 PyTuple_SET_ITEM(result, i, elem);
2614 Py_DECREF(oldelem);
2615 }
2616 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002617
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002618 Py_INCREF(result);
2619 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002620
2621empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002622 co->stopped = 1;
2623 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002624}
2625
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002626static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302627combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002628{
2629 if (lz->result == NULL) {
2630 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2631 } else if (lz->stopped) {
2632 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2633 } else {
2634 PyObject *indices;
2635 Py_ssize_t i;
2636
2637 /* we must pickle the indices and use them for setstate */
2638 indices = PyTuple_New(lz->r);
2639 if (!indices)
2640 return NULL;
2641 for (i=0; i<lz->r; i++)
2642 {
2643 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2644 if (!index) {
2645 Py_DECREF(indices);
2646 return NULL;
2647 }
2648 PyTuple_SET_ITEM(indices, i, index);
2649 }
2650
2651 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2652 }
2653}
2654
2655static PyObject *
2656combinations_setstate(combinationsobject *lz, PyObject *state)
2657{
2658 PyObject *result;
2659 Py_ssize_t i;
2660 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2661
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002662 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002663 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2664 return NULL;
2665 }
2666
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002667 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002668 Py_ssize_t max;
2669 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2670 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002671
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002672 if (index == -1 && PyErr_Occurred())
2673 return NULL; /* not an integer */
2674 max = i + n - lz->r;
2675 /* clamp the index (beware of negative max) */
2676 if (index > max)
2677 index = max;
2678 if (index < 0)
2679 index = 0;
2680 lz->indices[i] = index;
2681 }
2682
2683 result = PyTuple_New(lz->r);
2684 if (result == NULL)
2685 return NULL;
2686 for (i=0; i<lz->r; i++) {
2687 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2688 Py_INCREF(element);
2689 PyTuple_SET_ITEM(result, i, element);
2690 }
2691
Serhiy Storchaka48842712016-04-06 09:45:48 +03002692 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002693 Py_RETURN_NONE;
2694}
2695
2696static PyMethodDef combinations_methods[] = {
2697 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2698 reduce_doc},
2699 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2700 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002701 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2702 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002703 {NULL, NULL} /* sentinel */
2704};
2705
Christian Heimes380f7f22008-02-28 11:19:05 +00002706static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002707 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002708 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002709 sizeof(combinationsobject), /* tp_basicsize */
2710 0, /* tp_itemsize */
2711 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002712 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002713 0, /* tp_print */
2714 0, /* tp_getattr */
2715 0, /* tp_setattr */
2716 0, /* tp_reserved */
2717 0, /* tp_repr */
2718 0, /* tp_as_number */
2719 0, /* tp_as_sequence */
2720 0, /* tp_as_mapping */
2721 0, /* tp_hash */
2722 0, /* tp_call */
2723 0, /* tp_str */
2724 PyObject_GenericGetAttr, /* tp_getattro */
2725 0, /* tp_setattro */
2726 0, /* tp_as_buffer */
2727 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2728 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002729 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002730 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002731 0, /* tp_clear */
2732 0, /* tp_richcompare */
2733 0, /* tp_weaklistoffset */
2734 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002735 (iternextfunc)combinations_next, /* tp_iternext */
2736 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002737 0, /* tp_members */
2738 0, /* tp_getset */
2739 0, /* tp_base */
2740 0, /* tp_dict */
2741 0, /* tp_descr_get */
2742 0, /* tp_descr_set */
2743 0, /* tp_dictoffset */
2744 0, /* tp_init */
2745 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002746 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002747 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002748};
2749
2750
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002751/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002752
2753/* Equivalent to:
2754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002755 def combinations_with_replacement(iterable, r):
2756 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2757 # number items returned: (n+r-1)! / r! / (n-1)!
2758 pool = tuple(iterable)
2759 n = len(pool)
2760 indices = [0] * r
2761 yield tuple(pool[i] for i in indices)
2762 while 1:
2763 for i in reversed(range(r)):
2764 if indices[i] != n - 1:
2765 break
2766 else:
2767 return
2768 indices[i:] = [indices[i] + 1] * (r - i)
2769 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002771 def combinations_with_replacement2(iterable, r):
2772 'Alternate version that filters from product()'
2773 pool = tuple(iterable)
2774 n = len(pool)
2775 for indices in product(range(n), repeat=r):
2776 if sorted(indices) == list(indices):
2777 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002778*/
2779typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002780 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002781 PyObject *pool; /* input converted to a tuple */
2782 Py_ssize_t *indices; /* one index per result element */
2783 PyObject *result; /* most recently returned result tuple */
2784 Py_ssize_t r; /* size of result tuple */
2785 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002786} cwrobject;
2787
2788static PyTypeObject cwr_type;
2789
Tal Einatc4bccd32018-09-12 00:49:13 +03002790/*[clinic input]
2791@classmethod
2792itertools.combinations_with_replacement.__new__
2793 iterable: object
2794 r: Py_ssize_t
2795Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2796
2797combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2798[clinic start generated code]*/
2799
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002800static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002801itertools_combinations_with_replacement_impl(PyTypeObject *type,
2802 PyObject *iterable,
2803 Py_ssize_t r)
2804/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002805{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002806 cwrobject *co;
2807 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002808 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 Py_ssize_t *indices = NULL;
2810 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 pool = PySequence_Tuple(iterable);
2813 if (pool == NULL)
2814 goto error;
2815 n = PyTuple_GET_SIZE(pool);
2816 if (r < 0) {
2817 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2818 goto error;
2819 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002820
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002821 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 if (indices == NULL) {
2823 PyErr_NoMemory();
2824 goto error;
2825 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002826
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002827 for (i=0 ; i<r ; i++)
2828 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002829
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 /* create cwrobject structure */
2831 co = (cwrobject *)type->tp_alloc(type, 0);
2832 if (co == NULL)
2833 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002834
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002835 co->pool = pool;
2836 co->indices = indices;
2837 co->result = NULL;
2838 co->r = r;
2839 co->stopped = !n && r;
2840
2841 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002842
2843error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002844 if (indices != NULL)
2845 PyMem_Free(indices);
2846 Py_XDECREF(pool);
2847 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002848}
2849
2850static void
2851cwr_dealloc(cwrobject *co)
2852{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002853 PyObject_GC_UnTrack(co);
2854 Py_XDECREF(co->pool);
2855 Py_XDECREF(co->result);
2856 if (co->indices != NULL)
2857 PyMem_Free(co->indices);
2858 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002859}
2860
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002861static PyObject *
2862cwr_sizeof(cwrobject *co, void *unused)
2863{
2864 Py_ssize_t res;
2865
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002866 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002867 res += co->r * sizeof(Py_ssize_t);
2868 return PyLong_FromSsize_t(res);
2869}
2870
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002871static int
2872cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002874 Py_VISIT(co->pool);
2875 Py_VISIT(co->result);
2876 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002877}
2878
2879static PyObject *
2880cwr_next(cwrobject *co)
2881{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002882 PyObject *elem;
2883 PyObject *oldelem;
2884 PyObject *pool = co->pool;
2885 Py_ssize_t *indices = co->indices;
2886 PyObject *result = co->result;
2887 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2888 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002889 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002891 if (co->stopped)
2892 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002893
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002894 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002895 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002896 result = PyTuple_New(r);
2897 if (result == NULL)
2898 goto empty;
2899 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002900 if (n > 0) {
2901 elem = PyTuple_GET_ITEM(pool, 0);
2902 for (i=0; i<r ; i++) {
2903 assert(indices[i] == 0);
2904 Py_INCREF(elem);
2905 PyTuple_SET_ITEM(result, i, elem);
2906 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002907 }
2908 } else {
2909 /* Copy the previous result tuple or re-use it if available */
2910 if (Py_REFCNT(result) > 1) {
2911 PyObject *old_result = result;
2912 result = PyTuple_New(r);
2913 if (result == NULL)
2914 goto empty;
2915 co->result = result;
2916 for (i=0; i<r ; i++) {
2917 elem = PyTuple_GET_ITEM(old_result, i);
2918 Py_INCREF(elem);
2919 PyTuple_SET_ITEM(result, i, elem);
2920 }
2921 Py_DECREF(old_result);
2922 }
2923 /* Now, we've got the only copy so we can update it in-place CPython's
2924 empty tuple is a singleton and cached in PyTuple's freelist. */
2925 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002926
Tim Peters9edb1682013-09-03 11:49:31 -05002927 /* Scan indices right-to-left until finding one that is not
2928 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002929 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2930 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002931
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002932 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002933 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002934 if (i < 0)
2935 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002936
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002937 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002938 maximum. Then set all to the right to the same value. */
2939 index = indices[i] + 1;
2940 assert(index < n);
2941 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002942 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002943 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002944 Py_INCREF(elem);
2945 oldelem = PyTuple_GET_ITEM(result, i);
2946 PyTuple_SET_ITEM(result, i, elem);
2947 Py_DECREF(oldelem);
2948 }
2949 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002950
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002951 Py_INCREF(result);
2952 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002953
2954empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002955 co->stopped = 1;
2956 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002957}
2958
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002959static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302960cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002961{
2962 if (lz->result == NULL) {
2963 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2964 } else if (lz->stopped) {
2965 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2966 } else {
2967 PyObject *indices;
2968 Py_ssize_t i;
2969
2970 /* we must pickle the indices and use them for setstate */
2971 indices = PyTuple_New(lz->r);
2972 if (!indices)
2973 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002974 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002975 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2976 if (!index) {
2977 Py_DECREF(indices);
2978 return NULL;
2979 }
2980 PyTuple_SET_ITEM(indices, i, index);
2981 }
2982
2983 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2984 }
2985}
2986
2987static PyObject *
2988cwr_setstate(cwrobject *lz, PyObject *state)
2989{
2990 PyObject *result;
2991 Py_ssize_t n, i;
2992
2993 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2994 {
2995 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2996 return NULL;
2997 }
2998
2999 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003000 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003001 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
3002 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003003
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003004 if (index < 0 && PyErr_Occurred())
3005 return NULL; /* not an integer */
3006 /* clamp the index */
3007 if (index < 0)
3008 index = 0;
3009 else if (index > n-1)
3010 index = n-1;
3011 lz->indices[i] = index;
3012 }
3013 result = PyTuple_New(lz->r);
3014 if (result == NULL)
3015 return NULL;
3016 for (i=0; i<lz->r; i++) {
3017 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3018 Py_INCREF(element);
3019 PyTuple_SET_ITEM(result, i, element);
3020 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003021 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003022 Py_RETURN_NONE;
3023}
3024
3025static PyMethodDef cwr_methods[] = {
3026 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3027 reduce_doc},
3028 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3029 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003030 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3031 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003032 {NULL, NULL} /* sentinel */
3033};
3034
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003035static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003036 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003037 "itertools.combinations_with_replacement", /* tp_name */
3038 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003039 0, /* tp_itemsize */
3040 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003041 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003042 0, /* tp_print */
3043 0, /* tp_getattr */
3044 0, /* tp_setattr */
3045 0, /* tp_reserved */
3046 0, /* tp_repr */
3047 0, /* tp_as_number */
3048 0, /* tp_as_sequence */
3049 0, /* tp_as_mapping */
3050 0, /* tp_hash */
3051 0, /* tp_call */
3052 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003053 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003054 0, /* tp_setattro */
3055 0, /* tp_as_buffer */
3056 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003057 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003058 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003059 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003060 0, /* tp_clear */
3061 0, /* tp_richcompare */
3062 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003063 PyObject_SelfIter, /* tp_iter */
3064 (iternextfunc)cwr_next, /* tp_iternext */
3065 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 0, /* tp_members */
3067 0, /* tp_getset */
3068 0, /* tp_base */
3069 0, /* tp_dict */
3070 0, /* tp_descr_get */
3071 0, /* tp_descr_set */
3072 0, /* tp_dictoffset */
3073 0, /* tp_init */
3074 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003075 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003076 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003077};
3078
3079
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003080/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003081
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003082def permutations(iterable, r=None):
3083 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3084 pool = tuple(iterable)
3085 n = len(pool)
3086 r = n if r is None else r
3087 indices = range(n)
3088 cycles = range(n-r+1, n+1)[::-1]
3089 yield tuple(pool[i] for i in indices[:r])
3090 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003091 for i in reversed(range(r)):
3092 cycles[i] -= 1
3093 if cycles[i] == 0:
3094 indices[i:] = indices[i+1:] + indices[i:i+1]
3095 cycles[i] = n - i
3096 else:
3097 j = cycles[i]
3098 indices[i], indices[-j] = indices[-j], indices[i]
3099 yield tuple(pool[i] for i in indices[:r])
3100 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003101 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003102 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003103*/
3104
3105typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003106 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003107 PyObject *pool; /* input converted to a tuple */
3108 Py_ssize_t *indices; /* one index per element in the pool */
3109 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3110 PyObject *result; /* most recently returned result tuple */
3111 Py_ssize_t r; /* size of result tuple */
3112 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003113} permutationsobject;
3114
3115static PyTypeObject permutations_type;
3116
Tal Einatc4bccd32018-09-12 00:49:13 +03003117/*[clinic input]
3118@classmethod
3119itertools.permutations.__new__
3120 iterable: object
3121 r as robj: object = None
3122Return successive r-length permutations of elements in the iterable.
3123
3124permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3125[clinic start generated code]*/
3126
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003127static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003128itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3129 PyObject *robj)
3130/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003131{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003132 permutationsobject *po;
3133 Py_ssize_t n;
3134 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003135 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003136 Py_ssize_t *indices = NULL;
3137 Py_ssize_t *cycles = NULL;
3138 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003139
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003140 pool = PySequence_Tuple(iterable);
3141 if (pool == NULL)
3142 goto error;
3143 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003144
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003145 r = n;
3146 if (robj != Py_None) {
3147 if (!PyLong_Check(robj)) {
3148 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3149 goto error;
3150 }
3151 r = PyLong_AsSsize_t(robj);
3152 if (r == -1 && PyErr_Occurred())
3153 goto error;
3154 }
3155 if (r < 0) {
3156 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3157 goto error;
3158 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003159
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003160 indices = PyMem_New(Py_ssize_t, n);
3161 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003162 if (indices == NULL || cycles == NULL) {
3163 PyErr_NoMemory();
3164 goto error;
3165 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003167 for (i=0 ; i<n ; i++)
3168 indices[i] = i;
3169 for (i=0 ; i<r ; i++)
3170 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003171
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003172 /* create permutationsobject structure */
3173 po = (permutationsobject *)type->tp_alloc(type, 0);
3174 if (po == NULL)
3175 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003176
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003177 po->pool = pool;
3178 po->indices = indices;
3179 po->cycles = cycles;
3180 po->result = NULL;
3181 po->r = r;
3182 po->stopped = r > n ? 1 : 0;
3183
3184 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003185
3186error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003187 if (indices != NULL)
3188 PyMem_Free(indices);
3189 if (cycles != NULL)
3190 PyMem_Free(cycles);
3191 Py_XDECREF(pool);
3192 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003193}
3194
3195static void
3196permutations_dealloc(permutationsobject *po)
3197{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003198 PyObject_GC_UnTrack(po);
3199 Py_XDECREF(po->pool);
3200 Py_XDECREF(po->result);
3201 PyMem_Free(po->indices);
3202 PyMem_Free(po->cycles);
3203 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003204}
3205
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003206static PyObject *
3207permutations_sizeof(permutationsobject *po, void *unused)
3208{
3209 Py_ssize_t res;
3210
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003211 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003212 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3213 res += po->r * sizeof(Py_ssize_t);
3214 return PyLong_FromSsize_t(res);
3215}
3216
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003217static int
3218permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3219{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003220 Py_VISIT(po->pool);
3221 Py_VISIT(po->result);
3222 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003223}
3224
3225static PyObject *
3226permutations_next(permutationsobject *po)
3227{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003228 PyObject *elem;
3229 PyObject *oldelem;
3230 PyObject *pool = po->pool;
3231 Py_ssize_t *indices = po->indices;
3232 Py_ssize_t *cycles = po->cycles;
3233 PyObject *result = po->result;
3234 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3235 Py_ssize_t r = po->r;
3236 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003237
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003238 if (po->stopped)
3239 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003240
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003241 if (result == NULL) {
3242 /* On the first pass, initialize result tuple using the indices */
3243 result = PyTuple_New(r);
3244 if (result == NULL)
3245 goto empty;
3246 po->result = result;
3247 for (i=0; i<r ; i++) {
3248 index = indices[i];
3249 elem = PyTuple_GET_ITEM(pool, index);
3250 Py_INCREF(elem);
3251 PyTuple_SET_ITEM(result, i, elem);
3252 }
3253 } else {
3254 if (n == 0)
3255 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003256
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003257 /* Copy the previous result tuple or re-use it if available */
3258 if (Py_REFCNT(result) > 1) {
3259 PyObject *old_result = result;
3260 result = PyTuple_New(r);
3261 if (result == NULL)
3262 goto empty;
3263 po->result = result;
3264 for (i=0; i<r ; i++) {
3265 elem = PyTuple_GET_ITEM(old_result, i);
3266 Py_INCREF(elem);
3267 PyTuple_SET_ITEM(result, i, elem);
3268 }
3269 Py_DECREF(old_result);
3270 }
3271 /* Now, we've got the only copy so we can update it in-place */
3272 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003273
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003274 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3275 for (i=r-1 ; i>=0 ; i--) {
3276 cycles[i] -= 1;
3277 if (cycles[i] == 0) {
3278 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3279 index = indices[i];
3280 for (j=i ; j<n-1 ; j++)
3281 indices[j] = indices[j+1];
3282 indices[n-1] = index;
3283 cycles[i] = n - i;
3284 } else {
3285 j = cycles[i];
3286 index = indices[i];
3287 indices[i] = indices[n-j];
3288 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003290 for (k=i; k<r ; k++) {
3291 /* start with i, the leftmost element that changed */
3292 /* yield tuple(pool[k] for k in indices[:r]) */
3293 index = indices[k];
3294 elem = PyTuple_GET_ITEM(pool, index);
3295 Py_INCREF(elem);
3296 oldelem = PyTuple_GET_ITEM(result, k);
3297 PyTuple_SET_ITEM(result, k, elem);
3298 Py_DECREF(oldelem);
3299 }
3300 break;
3301 }
3302 }
3303 /* If i is negative, then the cycles have all
3304 rolled-over and we're done. */
3305 if (i < 0)
3306 goto empty;
3307 }
3308 Py_INCREF(result);
3309 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003310
3311empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003312 po->stopped = 1;
3313 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003314}
3315
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003316static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303317permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003318{
3319 if (po->result == NULL) {
3320 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3321 } else if (po->stopped) {
3322 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3323 } else {
3324 PyObject *indices=NULL, *cycles=NULL;
3325 Py_ssize_t n, i;
3326
3327 /* we must pickle the indices and cycles and use them for setstate */
3328 n = PyTuple_GET_SIZE(po->pool);
3329 indices = PyTuple_New(n);
3330 if (indices == NULL)
3331 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003332 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003333 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3334 if (!index)
3335 goto err;
3336 PyTuple_SET_ITEM(indices, i, index);
3337 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003338
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003339 cycles = PyTuple_New(po->r);
3340 if (cycles == NULL)
3341 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003342 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003343 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3344 if (!index)
3345 goto err;
3346 PyTuple_SET_ITEM(cycles, i, index);
3347 }
3348 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3349 po->pool, po->r,
3350 indices, cycles);
3351 err:
3352 Py_XDECREF(indices);
3353 Py_XDECREF(cycles);
3354 return NULL;
3355 }
3356}
3357
3358static PyObject *
3359permutations_setstate(permutationsobject *po, PyObject *state)
3360{
3361 PyObject *indices, *cycles, *result;
3362 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003363
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003364 if (!PyTuple_Check(state)) {
3365 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3366 return NULL;
3367 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003368 if (!PyArg_ParseTuple(state, "O!O!",
3369 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003370 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003371 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003372 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003373
3374 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003375 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003376 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3377 return NULL;
3378 }
3379
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003380 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003381 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3382 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3383 if (index < 0 && PyErr_Occurred())
3384 return NULL; /* not an integer */
3385 /* clamp the index */
3386 if (index < 0)
3387 index = 0;
3388 else if (index > n-1)
3389 index = n-1;
3390 po->indices[i] = index;
3391 }
3392
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003393 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003394 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3395 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3396 if (index < 0 && PyErr_Occurred())
3397 return NULL; /* not an integer */
3398 if (index < 1)
3399 index = 1;
3400 else if (index > n-i)
3401 index = n-i;
3402 po->cycles[i] = index;
3403 }
3404 result = PyTuple_New(po->r);
3405 if (result == NULL)
3406 return NULL;
3407 for (i=0; i<po->r; i++) {
3408 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3409 Py_INCREF(element);
3410 PyTuple_SET_ITEM(result, i, element);
3411 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003412 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003413 Py_RETURN_NONE;
3414}
3415
3416static PyMethodDef permuations_methods[] = {
3417 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3418 reduce_doc},
3419 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3420 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003421 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3422 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003423 {NULL, NULL} /* sentinel */
3424};
3425
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003426static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003427 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003428 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003429 sizeof(permutationsobject), /* tp_basicsize */
3430 0, /* tp_itemsize */
3431 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003432 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003433 0, /* tp_print */
3434 0, /* tp_getattr */
3435 0, /* tp_setattr */
3436 0, /* tp_reserved */
3437 0, /* tp_repr */
3438 0, /* tp_as_number */
3439 0, /* tp_as_sequence */
3440 0, /* tp_as_mapping */
3441 0, /* tp_hash */
3442 0, /* tp_call */
3443 0, /* tp_str */
3444 PyObject_GenericGetAttr, /* tp_getattro */
3445 0, /* tp_setattro */
3446 0, /* tp_as_buffer */
3447 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3448 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003449 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003450 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003451 0, /* tp_clear */
3452 0, /* tp_richcompare */
3453 0, /* tp_weaklistoffset */
3454 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003455 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003456 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003457 0, /* tp_members */
3458 0, /* tp_getset */
3459 0, /* tp_base */
3460 0, /* tp_dict */
3461 0, /* tp_descr_get */
3462 0, /* tp_descr_set */
3463 0, /* tp_dictoffset */
3464 0, /* tp_init */
3465 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003466 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003467 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003468};
3469
Tal Einatc4bccd32018-09-12 00:49:13 +03003470
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003471/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003472
3473typedef struct {
3474 PyObject_HEAD
3475 PyObject *total;
3476 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003477 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003478} accumulateobject;
3479
3480static PyTypeObject accumulate_type;
3481
Tal Einatc4bccd32018-09-12 00:49:13 +03003482/*[clinic input]
3483@classmethod
3484itertools.accumulate.__new__
3485 iterable: object
3486 func as binop: object = None
3487Return series of accumulated sums (or other binary function results).
3488[clinic start generated code]*/
3489
Raymond Hettinger482ba772010-12-01 22:48:00 +00003490static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003491itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
3492 PyObject *binop)
3493/*[clinic end generated code: output=514d0fb30ba14d55 input=6d9d16aaa1d3cbfc]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003494{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003495 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003496 accumulateobject *lz;
3497
Raymond Hettinger482ba772010-12-01 22:48:00 +00003498
3499 /* Get iterator. */
3500 it = PyObject_GetIter(iterable);
3501 if (it == NULL)
3502 return NULL;
3503
Raymond Hettinger482ba772010-12-01 22:48:00 +00003504 /* create accumulateobject structure */
3505 lz = (accumulateobject *)type->tp_alloc(type, 0);
3506 if (lz == NULL) {
3507 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003508 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003509 }
3510
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003511 if (binop != Py_None) {
3512 Py_XINCREF(binop);
3513 lz->binop = binop;
3514 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003515 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003516 lz->it = it;
3517 return (PyObject *)lz;
3518}
3519
3520static void
3521accumulate_dealloc(accumulateobject *lz)
3522{
3523 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003524 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003525 Py_XDECREF(lz->total);
3526 Py_XDECREF(lz->it);
3527 Py_TYPE(lz)->tp_free(lz);
3528}
3529
3530static int
3531accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3532{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003533 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003534 Py_VISIT(lz->it);
3535 Py_VISIT(lz->total);
3536 return 0;
3537}
3538
3539static PyObject *
3540accumulate_next(accumulateobject *lz)
3541{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003542 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003543
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003544 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003545 if (val == NULL)
3546 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003547
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003548 if (lz->total == NULL) {
3549 Py_INCREF(val);
3550 lz->total = val;
3551 return lz->total;
3552 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003553
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003554 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003555 newtotal = PyNumber_Add(lz->total, val);
3556 else
3557 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003558 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003559 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003560 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003561
Raymond Hettinger482ba772010-12-01 22:48:00 +00003562 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003563 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003564 return newtotal;
3565}
3566
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003567static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303568accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003569{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003570 if (lz->total == Py_None) {
3571 PyObject *it;
3572
3573 if (PyType_Ready(&chain_type) < 0)
3574 return NULL;
3575 if (PyType_Ready(&islice_type) < 0)
3576 return NULL;
3577 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3578 lz->total, lz->it);
3579 if (it == NULL)
3580 return NULL;
3581 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3582 it, lz->binop ? lz->binop : Py_None);
3583 if (it == NULL)
3584 return NULL;
3585 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3586 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003587 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3588 lz->it, lz->binop?lz->binop:Py_None,
3589 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003590}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003591
3592static PyObject *
3593accumulate_setstate(accumulateobject *lz, PyObject *state)
3594{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003595 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003596 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003597 Py_RETURN_NONE;
3598}
3599
3600static PyMethodDef accumulate_methods[] = {
3601 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3602 reduce_doc},
3603 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3604 setstate_doc},
3605 {NULL, NULL} /* sentinel */
3606};
3607
Raymond Hettinger482ba772010-12-01 22:48:00 +00003608static PyTypeObject accumulate_type = {
3609 PyVarObject_HEAD_INIT(NULL, 0)
3610 "itertools.accumulate", /* tp_name */
3611 sizeof(accumulateobject), /* tp_basicsize */
3612 0, /* tp_itemsize */
3613 /* methods */
3614 (destructor)accumulate_dealloc, /* tp_dealloc */
3615 0, /* tp_print */
3616 0, /* tp_getattr */
3617 0, /* tp_setattr */
3618 0, /* tp_reserved */
3619 0, /* tp_repr */
3620 0, /* tp_as_number */
3621 0, /* tp_as_sequence */
3622 0, /* tp_as_mapping */
3623 0, /* tp_hash */
3624 0, /* tp_call */
3625 0, /* tp_str */
3626 PyObject_GenericGetAttr, /* tp_getattro */
3627 0, /* tp_setattro */
3628 0, /* tp_as_buffer */
3629 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3630 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003631 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003632 (traverseproc)accumulate_traverse, /* tp_traverse */
3633 0, /* tp_clear */
3634 0, /* tp_richcompare */
3635 0, /* tp_weaklistoffset */
3636 PyObject_SelfIter, /* tp_iter */
3637 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003638 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003639 0, /* tp_members */
3640 0, /* tp_getset */
3641 0, /* tp_base */
3642 0, /* tp_dict */
3643 0, /* tp_descr_get */
3644 0, /* tp_descr_set */
3645 0, /* tp_dictoffset */
3646 0, /* tp_init */
3647 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003648 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003649 PyObject_GC_Del, /* tp_free */
3650};
3651
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003652
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003653/* compress object ************************************************************/
3654
3655/* Equivalent to:
3656
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003657 def compress(data, selectors):
3658 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3659 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003660*/
3661
3662typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003663 PyObject_HEAD
3664 PyObject *data;
3665 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003666} compressobject;
3667
3668static PyTypeObject compress_type;
3669
Tal Einatc4bccd32018-09-12 00:49:13 +03003670/*[clinic input]
3671@classmethod
3672itertools.compress.__new__
3673 data as seq1: object
3674 selectors as seq2: object
3675Return data elements corresponding to true selector elements.
3676
3677Forms a shorter iterator from selected data elements using the selectors to
3678choose the data elements.
3679[clinic start generated code]*/
3680
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003681static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003682itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3683/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003685 PyObject *data=NULL, *selectors=NULL;
3686 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003688 data = PyObject_GetIter(seq1);
3689 if (data == NULL)
3690 goto fail;
3691 selectors = PyObject_GetIter(seq2);
3692 if (selectors == NULL)
3693 goto fail;
3694
3695 /* create compressobject structure */
3696 lz = (compressobject *)type->tp_alloc(type, 0);
3697 if (lz == NULL)
3698 goto fail;
3699 lz->data = data;
3700 lz->selectors = selectors;
3701 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003702
3703fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 Py_XDECREF(data);
3705 Py_XDECREF(selectors);
3706 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003707}
3708
3709static void
3710compress_dealloc(compressobject *lz)
3711{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003712 PyObject_GC_UnTrack(lz);
3713 Py_XDECREF(lz->data);
3714 Py_XDECREF(lz->selectors);
3715 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003716}
3717
3718static int
3719compress_traverse(compressobject *lz, visitproc visit, void *arg)
3720{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003721 Py_VISIT(lz->data);
3722 Py_VISIT(lz->selectors);
3723 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003724}
3725
3726static PyObject *
3727compress_next(compressobject *lz)
3728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003729 PyObject *data = lz->data, *selectors = lz->selectors;
3730 PyObject *datum, *selector;
3731 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3732 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3733 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003735 while (1) {
3736 /* Steps: get datum, get selector, evaluate selector.
3737 Order is important (to match the pure python version
3738 in terms of which input gets a chance to raise an
3739 exception first).
3740 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003741
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003742 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003743 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003744 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003746 selector = selectornext(selectors);
3747 if (selector == NULL) {
3748 Py_DECREF(datum);
3749 return NULL;
3750 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003752 ok = PyObject_IsTrue(selector);
3753 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003754 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003755 return datum;
3756 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003757 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003758 return NULL;
3759 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003760}
3761
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003762static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303763compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003764{
3765 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3766 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003767}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003768
3769static PyMethodDef compress_methods[] = {
3770 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3771 reduce_doc},
3772 {NULL, NULL} /* sentinel */
3773};
3774
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003775static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003776 PyVarObject_HEAD_INIT(NULL, 0)
3777 "itertools.compress", /* tp_name */
3778 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003779 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003780 /* methods */
3781 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003782 0, /* tp_print */
3783 0, /* tp_getattr */
3784 0, /* tp_setattr */
3785 0, /* tp_reserved */
3786 0, /* tp_repr */
3787 0, /* tp_as_number */
3788 0, /* tp_as_sequence */
3789 0, /* tp_as_mapping */
3790 0, /* tp_hash */
3791 0, /* tp_call */
3792 0, /* tp_str */
3793 PyObject_GenericGetAttr, /* tp_getattro */
3794 0, /* tp_setattro */
3795 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003796 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003797 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003798 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003799 (traverseproc)compress_traverse, /* tp_traverse */
3800 0, /* tp_clear */
3801 0, /* tp_richcompare */
3802 0, /* tp_weaklistoffset */
3803 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003804 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003805 compress_methods, /* tp_methods */
3806 0, /* tp_members */
3807 0, /* tp_getset */
3808 0, /* tp_base */
3809 0, /* tp_dict */
3810 0, /* tp_descr_get */
3811 0, /* tp_descr_set */
3812 0, /* tp_dictoffset */
3813 0, /* tp_init */
3814 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003815 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003816 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003817};
3818
3819
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003820/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003821
3822typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003823 PyObject_HEAD
3824 PyObject *func;
3825 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003826} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003827
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003828static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003829
Tal Einatc4bccd32018-09-12 00:49:13 +03003830/*[clinic input]
3831@classmethod
3832itertools.filterfalse.__new__
3833 function as func: object
3834 iterable as seq: object
3835 /
3836Return those items of iterable for which function(item) is false.
3837
3838If function is None, return the items that are false.
3839[clinic start generated code]*/
3840
Raymond Hettinger60eca932003-02-09 06:40:58 +00003841static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003842itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3843/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003844{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003845 PyObject *it;
3846 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003847
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003848 /* Get iterator. */
3849 it = PyObject_GetIter(seq);
3850 if (it == NULL)
3851 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003852
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003853 /* create filterfalseobject structure */
3854 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3855 if (lz == NULL) {
3856 Py_DECREF(it);
3857 return NULL;
3858 }
3859 Py_INCREF(func);
3860 lz->func = func;
3861 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003862
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003863 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003864}
3865
3866static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003867filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003869 PyObject_GC_UnTrack(lz);
3870 Py_XDECREF(lz->func);
3871 Py_XDECREF(lz->it);
3872 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003873}
3874
3875static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003876filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003877{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003878 Py_VISIT(lz->it);
3879 Py_VISIT(lz->func);
3880 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003881}
3882
3883static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003884filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003885{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003886 PyObject *item;
3887 PyObject *it = lz->it;
3888 long ok;
3889 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003891 iternext = *Py_TYPE(it)->tp_iternext;
3892 for (;;) {
3893 item = iternext(it);
3894 if (item == NULL)
3895 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003896
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003897 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3898 ok = PyObject_IsTrue(item);
3899 } else {
3900 PyObject *good;
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01003901 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003902 if (good == NULL) {
3903 Py_DECREF(item);
3904 return NULL;
3905 }
3906 ok = PyObject_IsTrue(good);
3907 Py_DECREF(good);
3908 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003909 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003910 return item;
3911 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003912 if (ok < 0)
3913 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003914 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003915}
3916
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003917static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303918filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003919{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003920 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3921}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003922
3923static PyMethodDef filterfalse_methods[] = {
3924 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3925 reduce_doc},
3926 {NULL, NULL} /* sentinel */
3927};
3928
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003929static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003930 PyVarObject_HEAD_INIT(NULL, 0)
3931 "itertools.filterfalse", /* tp_name */
3932 sizeof(filterfalseobject), /* tp_basicsize */
3933 0, /* tp_itemsize */
3934 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003935 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003936 0, /* tp_print */
3937 0, /* tp_getattr */
3938 0, /* tp_setattr */
3939 0, /* tp_reserved */
3940 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 */
3950 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3951 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003952 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003953 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003954 0, /* tp_clear */
3955 0, /* tp_richcompare */
3956 0, /* tp_weaklistoffset */
3957 PyObject_SelfIter, /* tp_iter */
3958 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003959 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003960 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_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003970 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003971};
3972
3973
3974/* count object ************************************************************/
3975
3976typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003977 PyObject_HEAD
3978 Py_ssize_t cnt;
3979 PyObject *long_cnt;
3980 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003981} countobject;
3982
Raymond Hettinger30729212009-02-12 06:28:27 +00003983/* Counting logic and invariants:
3984
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003985fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003986
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003987 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003988 Advances with: cnt += 1
3989 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003990
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003991slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003992
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003993 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3994 All counting is done with python objects (no overflows or underflows).
3995 Advances with: long_cnt += long_step
3996 Step may be zero -- effectively a slow version of repeat(cnt).
3997 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003998*/
3999
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004000static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004001
Tal Einatc4bccd32018-09-12 00:49:13 +03004002/*[clinic input]
4003@classmethod
4004itertools.count.__new__
4005 start as long_cnt: object(c_default="NULL") = 0
4006 step as long_step: object(c_default="NULL") = 1
4007Return a count object whose .__next__() method returns consecutive values.
4008
4009Equivalent to:
4010 def count(firstval=0, step=1):
4011 x = firstval
4012 while 1:
4013 yield x
4014 x += step
4015[clinic start generated code]*/
4016
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004018itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4019 PyObject *long_step)
4020/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004021{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004022 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004023 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004024 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004025 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004026
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004027 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4028 (long_step != NULL && !PyNumber_Check(long_step))) {
4029 PyErr_SetString(PyExc_TypeError, "a number is required");
4030 return NULL;
4031 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004032
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004033 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4034 (long_step == NULL || PyLong_Check(long_step));
4035
4036 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004037 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004038 if (fast_mode) {
4039 assert(PyLong_Check(long_cnt));
4040 cnt = PyLong_AsSsize_t(long_cnt);
4041 if (cnt == -1 && PyErr_Occurred()) {
4042 PyErr_Clear();
4043 fast_mode = 0;
4044 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004045 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004046 } else {
4047 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004048 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004049 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004050 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004051
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004052 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004053 if (long_step == NULL)
4054 long_step = _PyLong_One;
4055 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004057 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004058
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004059 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004060 if (fast_mode) {
4061 assert(PyLong_Check(long_step));
4062 step = PyLong_AsLong(long_step);
4063 if (step != 1) {
4064 fast_mode = 0;
4065 if (step == -1 && PyErr_Occurred())
4066 PyErr_Clear();
4067 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004068 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004069
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004070 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004071 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004072 else
4073 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004074
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004075 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4076 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4077 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004078 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004079
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004080 /* create countobject structure */
4081 lz = (countobject *)type->tp_alloc(type, 0);
4082 if (lz == NULL) {
4083 Py_XDECREF(long_cnt);
4084 return NULL;
4085 }
4086 lz->cnt = cnt;
4087 lz->long_cnt = long_cnt;
4088 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004090 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004091}
4092
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004093static void
4094count_dealloc(countobject *lz)
4095{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004096 PyObject_GC_UnTrack(lz);
4097 Py_XDECREF(lz->long_cnt);
4098 Py_XDECREF(lz->long_step);
4099 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004100}
4101
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004102static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004103count_traverse(countobject *lz, visitproc visit, void *arg)
4104{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 Py_VISIT(lz->long_cnt);
4106 Py_VISIT(lz->long_step);
4107 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004108}
4109
4110static PyObject *
4111count_nextlong(countobject *lz)
4112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004113 PyObject *long_cnt;
4114 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004115
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004116 long_cnt = lz->long_cnt;
4117 if (long_cnt == NULL) {
4118 /* Switch to slow_mode */
4119 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4120 if (long_cnt == NULL)
4121 return NULL;
4122 }
4123 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004124
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004125 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4126 if (stepped_up == NULL)
4127 return NULL;
4128 lz->long_cnt = stepped_up;
4129 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004130}
4131
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004132static PyObject *
4133count_next(countobject *lz)
4134{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004135 if (lz->cnt == PY_SSIZE_T_MAX)
4136 return count_nextlong(lz);
4137 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004138}
4139
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004140static PyObject *
4141count_repr(countobject *lz)
4142{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004143 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004144 return PyUnicode_FromFormat("%s(%zd)",
4145 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004146
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004147 if (PyLong_Check(lz->long_step)) {
4148 long step = PyLong_AsLong(lz->long_step);
4149 if (step == -1 && PyErr_Occurred()) {
4150 PyErr_Clear();
4151 }
4152 if (step == 1) {
4153 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004154 return PyUnicode_FromFormat("%s(%R)",
4155 _PyType_Name(Py_TYPE(lz)),
4156 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004157 }
4158 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004159 return PyUnicode_FromFormat("%s(%R, %R)",
4160 _PyType_Name(Py_TYPE(lz)),
4161 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004162}
4163
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004164static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304165count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004166{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004167 if (lz->cnt == PY_SSIZE_T_MAX)
4168 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4169 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004170}
4171
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004172static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004173 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004174 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004175 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004176};
4177
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004178static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004179 PyVarObject_HEAD_INIT(NULL, 0)
4180 "itertools.count", /* tp_name */
4181 sizeof(countobject), /* tp_basicsize */
4182 0, /* tp_itemsize */
4183 /* methods */
4184 (destructor)count_dealloc, /* tp_dealloc */
4185 0, /* tp_print */
4186 0, /* tp_getattr */
4187 0, /* tp_setattr */
4188 0, /* tp_reserved */
4189 (reprfunc)count_repr, /* tp_repr */
4190 0, /* tp_as_number */
4191 0, /* tp_as_sequence */
4192 0, /* tp_as_mapping */
4193 0, /* tp_hash */
4194 0, /* tp_call */
4195 0, /* tp_str */
4196 PyObject_GenericGetAttr, /* tp_getattro */
4197 0, /* tp_setattro */
4198 0, /* tp_as_buffer */
4199 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004200 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004201 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004202 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004203 0, /* tp_clear */
4204 0, /* tp_richcompare */
4205 0, /* tp_weaklistoffset */
4206 PyObject_SelfIter, /* tp_iter */
4207 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004208 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004209 0, /* tp_members */
4210 0, /* tp_getset */
4211 0, /* tp_base */
4212 0, /* tp_dict */
4213 0, /* tp_descr_get */
4214 0, /* tp_descr_set */
4215 0, /* tp_dictoffset */
4216 0, /* tp_init */
4217 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004218 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004219 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004220};
4221
4222
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004223/* repeat object ************************************************************/
4224
4225typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004226 PyObject_HEAD
4227 PyObject *element;
4228 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004229} repeatobject;
4230
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004231static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004232
4233static PyObject *
4234repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4235{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004236 repeatobject *ro;
4237 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004238 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004239 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004240
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004241 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4242 &element, &cnt))
4243 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004244
Raymond Hettinger97d35552014-06-24 21:36:58 -07004245 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004246 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004247 /* Does user supply times argument? */
4248 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004249 cnt = 0;
4250
4251 ro = (repeatobject *)type->tp_alloc(type, 0);
4252 if (ro == NULL)
4253 return NULL;
4254 Py_INCREF(element);
4255 ro->element = element;
4256 ro->cnt = cnt;
4257 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004258}
4259
4260static void
4261repeat_dealloc(repeatobject *ro)
4262{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004263 PyObject_GC_UnTrack(ro);
4264 Py_XDECREF(ro->element);
4265 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004266}
4267
4268static int
4269repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4270{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004271 Py_VISIT(ro->element);
4272 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004273}
4274
4275static PyObject *
4276repeat_next(repeatobject *ro)
4277{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004278 if (ro->cnt == 0)
4279 return NULL;
4280 if (ro->cnt > 0)
4281 ro->cnt--;
4282 Py_INCREF(ro->element);
4283 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004284}
4285
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004286static PyObject *
4287repeat_repr(repeatobject *ro)
4288{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004289 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004290 return PyUnicode_FromFormat("%s(%R)",
4291 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004292 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004293 return PyUnicode_FromFormat("%s(%R, %zd)",
4294 _PyType_Name(Py_TYPE(ro)), ro->element,
4295 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004296}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004297
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004298static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304299repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004300{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004301 if (ro->cnt == -1) {
4302 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4303 return NULL;
4304 }
4305 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004306}
4307
Armin Rigof5b3e362006-02-11 21:32:43 +00004308PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004309
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004310static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304311repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004312{
4313 /* unpickle this so that a new repeat iterator is constructed with an
4314 * object, then call __setstate__ on it to set cnt
4315 */
4316 if (ro->cnt >= 0)
4317 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4318 else
4319 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4320}
4321
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004322static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004323 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004324 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004325 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004326};
4327
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004328PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004329"repeat(object [,times]) -> create an iterator which returns the object\n\
4330for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004331endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004332
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004333static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004334 PyVarObject_HEAD_INIT(NULL, 0)
4335 "itertools.repeat", /* tp_name */
4336 sizeof(repeatobject), /* tp_basicsize */
4337 0, /* tp_itemsize */
4338 /* methods */
4339 (destructor)repeat_dealloc, /* tp_dealloc */
4340 0, /* tp_print */
4341 0, /* tp_getattr */
4342 0, /* tp_setattr */
4343 0, /* tp_reserved */
4344 (reprfunc)repeat_repr, /* tp_repr */
4345 0, /* tp_as_number */
4346 0, /* tp_as_sequence */
4347 0, /* tp_as_mapping */
4348 0, /* tp_hash */
4349 0, /* tp_call */
4350 0, /* tp_str */
4351 PyObject_GenericGetAttr, /* tp_getattro */
4352 0, /* tp_setattro */
4353 0, /* tp_as_buffer */
4354 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4355 Py_TPFLAGS_BASETYPE, /* tp_flags */
4356 repeat_doc, /* tp_doc */
4357 (traverseproc)repeat_traverse, /* tp_traverse */
4358 0, /* tp_clear */
4359 0, /* tp_richcompare */
4360 0, /* tp_weaklistoffset */
4361 PyObject_SelfIter, /* tp_iter */
4362 (iternextfunc)repeat_next, /* tp_iternext */
4363 repeat_methods, /* tp_methods */
4364 0, /* tp_members */
4365 0, /* tp_getset */
4366 0, /* tp_base */
4367 0, /* tp_dict */
4368 0, /* tp_descr_get */
4369 0, /* tp_descr_set */
4370 0, /* tp_dictoffset */
4371 0, /* tp_init */
4372 0, /* tp_alloc */
4373 repeat_new, /* tp_new */
4374 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004375};
4376
Tal Einatc4bccd32018-09-12 00:49:13 +03004377
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004378/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004379
4380typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004381 PyObject_HEAD
4382 Py_ssize_t tuplesize;
4383 Py_ssize_t numactive;
4384 PyObject *ittuple; /* tuple of iterators */
4385 PyObject *result;
4386 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004387} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004388
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004389static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004390
4391static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004392zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004393{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004394 ziplongestobject *lz;
4395 Py_ssize_t i;
4396 PyObject *ittuple; /* tuple of iterators */
4397 PyObject *result;
4398 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004399 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004400
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004401 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004402 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004403 if (fillvalue == NULL || PyDict_GET_SIZE(kwds) > 1) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004404 PyErr_SetString(PyExc_TypeError,
4405 "zip_longest() got an unexpected keyword argument");
4406 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004407 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004408 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004409
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004410 /* args must be a tuple */
4411 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004412 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004413
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004414 /* obtain iterators */
4415 ittuple = PyTuple_New(tuplesize);
4416 if (ittuple == NULL)
4417 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004418 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004419 PyObject *item = PyTuple_GET_ITEM(args, i);
4420 PyObject *it = PyObject_GetIter(item);
4421 if (it == NULL) {
4422 if (PyErr_ExceptionMatches(PyExc_TypeError))
4423 PyErr_Format(PyExc_TypeError,
4424 "zip_longest argument #%zd must support iteration",
4425 i+1);
4426 Py_DECREF(ittuple);
4427 return NULL;
4428 }
4429 PyTuple_SET_ITEM(ittuple, i, it);
4430 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004431
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004432 /* create a result holder */
4433 result = PyTuple_New(tuplesize);
4434 if (result == NULL) {
4435 Py_DECREF(ittuple);
4436 return NULL;
4437 }
4438 for (i=0 ; i < tuplesize ; i++) {
4439 Py_INCREF(Py_None);
4440 PyTuple_SET_ITEM(result, i, Py_None);
4441 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004443 /* create ziplongestobject structure */
4444 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4445 if (lz == NULL) {
4446 Py_DECREF(ittuple);
4447 Py_DECREF(result);
4448 return NULL;
4449 }
4450 lz->ittuple = ittuple;
4451 lz->tuplesize = tuplesize;
4452 lz->numactive = tuplesize;
4453 lz->result = result;
4454 Py_INCREF(fillvalue);
4455 lz->fillvalue = fillvalue;
4456 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004457}
4458
4459static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004460zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004461{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004462 PyObject_GC_UnTrack(lz);
4463 Py_XDECREF(lz->ittuple);
4464 Py_XDECREF(lz->result);
4465 Py_XDECREF(lz->fillvalue);
4466 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004467}
4468
4469static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004470zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004471{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004472 Py_VISIT(lz->ittuple);
4473 Py_VISIT(lz->result);
4474 Py_VISIT(lz->fillvalue);
4475 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004476}
4477
4478static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004479zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004481 Py_ssize_t i;
4482 Py_ssize_t tuplesize = lz->tuplesize;
4483 PyObject *result = lz->result;
4484 PyObject *it;
4485 PyObject *item;
4486 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004488 if (tuplesize == 0)
4489 return NULL;
4490 if (lz->numactive == 0)
4491 return NULL;
4492 if (Py_REFCNT(result) == 1) {
4493 Py_INCREF(result);
4494 for (i=0 ; i < tuplesize ; i++) {
4495 it = PyTuple_GET_ITEM(lz->ittuple, i);
4496 if (it == NULL) {
4497 Py_INCREF(lz->fillvalue);
4498 item = lz->fillvalue;
4499 } else {
4500 item = PyIter_Next(it);
4501 if (item == NULL) {
4502 lz->numactive -= 1;
4503 if (lz->numactive == 0 || PyErr_Occurred()) {
4504 lz->numactive = 0;
4505 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004506 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004507 } else {
4508 Py_INCREF(lz->fillvalue);
4509 item = lz->fillvalue;
4510 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4511 Py_DECREF(it);
4512 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004513 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004514 }
4515 olditem = PyTuple_GET_ITEM(result, i);
4516 PyTuple_SET_ITEM(result, i, item);
4517 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004518 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004519 } else {
4520 result = PyTuple_New(tuplesize);
4521 if (result == NULL)
4522 return NULL;
4523 for (i=0 ; i < tuplesize ; i++) {
4524 it = PyTuple_GET_ITEM(lz->ittuple, i);
4525 if (it == NULL) {
4526 Py_INCREF(lz->fillvalue);
4527 item = lz->fillvalue;
4528 } else {
4529 item = PyIter_Next(it);
4530 if (item == NULL) {
4531 lz->numactive -= 1;
4532 if (lz->numactive == 0 || PyErr_Occurred()) {
4533 lz->numactive = 0;
4534 Py_DECREF(result);
4535 return NULL;
4536 } else {
4537 Py_INCREF(lz->fillvalue);
4538 item = lz->fillvalue;
4539 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4540 Py_DECREF(it);
4541 }
4542 }
4543 }
4544 PyTuple_SET_ITEM(result, i, item);
4545 }
4546 }
4547 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004548}
4549
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004550static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304551zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004552{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004553
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004554 /* Create a new tuple with empty sequences where appropriate to pickle.
4555 * Then use setstate to set the fillvalue
4556 */
4557 int i;
4558 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004559
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004560 if (args == NULL)
4561 return NULL;
4562 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4563 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4564 if (elem == NULL) {
4565 elem = PyTuple_New(0);
4566 if (elem == NULL) {
4567 Py_DECREF(args);
4568 return NULL;
4569 }
4570 } else
4571 Py_INCREF(elem);
4572 PyTuple_SET_ITEM(args, i, elem);
4573 }
4574 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4575}
4576
4577static PyObject *
4578zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4579{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004580 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004581 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004582 Py_RETURN_NONE;
4583}
4584
4585static PyMethodDef zip_longest_methods[] = {
4586 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4587 reduce_doc},
4588 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4589 setstate_doc},
4590 {NULL, NULL} /* sentinel */
4591};
4592
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004593PyDoc_STRVAR(zip_longest_doc,
4594"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004595\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004596Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004597the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004598method continues until the longest iterable in the argument sequence\n\
4599is exhausted and then it raises StopIteration. When the shorter iterables\n\
4600are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4601defaults to None or can be specified by a keyword argument.\n\
4602");
4603
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004604static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004605 PyVarObject_HEAD_INIT(NULL, 0)
4606 "itertools.zip_longest", /* tp_name */
4607 sizeof(ziplongestobject), /* tp_basicsize */
4608 0, /* tp_itemsize */
4609 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004610 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004611 0, /* tp_print */
4612 0, /* tp_getattr */
4613 0, /* tp_setattr */
4614 0, /* tp_reserved */
4615 0, /* tp_repr */
4616 0, /* tp_as_number */
4617 0, /* tp_as_sequence */
4618 0, /* tp_as_mapping */
4619 0, /* tp_hash */
4620 0, /* tp_call */
4621 0, /* tp_str */
4622 PyObject_GenericGetAttr, /* tp_getattro */
4623 0, /* tp_setattro */
4624 0, /* tp_as_buffer */
4625 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4626 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004627 zip_longest_doc, /* tp_doc */
4628 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004629 0, /* tp_clear */
4630 0, /* tp_richcompare */
4631 0, /* tp_weaklistoffset */
4632 PyObject_SelfIter, /* tp_iter */
4633 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004634 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004635 0, /* tp_members */
4636 0, /* tp_getset */
4637 0, /* tp_base */
4638 0, /* tp_dict */
4639 0, /* tp_descr_get */
4640 0, /* tp_descr_set */
4641 0, /* tp_dictoffset */
4642 0, /* tp_init */
4643 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004644 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004645 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004646};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004647
Tal Einatc4bccd32018-09-12 00:49:13 +03004648
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004649/* module level code ********************************************************/
4650
4651PyDoc_STRVAR(module_doc,
4652"Functional tools for creating and using iterators.\n\
4653\n\
4654Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004655count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004656cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004657repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004658\n\
4659Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004660accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004661chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4662chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004663compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4664dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4665groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004666filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004667islice(seq, [start,] stop [, step]) --> elements from\n\
4668 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004669starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004670tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004671takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004672zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004673\n\
4674Combinatoric generators:\n\
4675product(p, q, ... [repeat=1]) --> cartesian product\n\
4676permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004677combinations(p, r)\n\
4678combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004679");
4680
4681
Raymond Hettingerad983e72003-11-12 14:32:26 +00004682static PyMethodDef module_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03004683 ITERTOOLS_TEE_METHODDEF
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004684 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004685};
4686
Martin v. Löwis1a214512008-06-11 05:26:20 +00004687
4688static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004689 PyModuleDef_HEAD_INIT,
4690 "itertools",
4691 module_doc,
4692 -1,
4693 module_methods,
4694 NULL,
4695 NULL,
4696 NULL,
4697 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004698};
4699
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004700PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004701PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004702{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004703 int i;
4704 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004705 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004706 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004707 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004708 &combinations_type,
4709 &cwr_type,
4710 &cycle_type,
4711 &dropwhile_type,
4712 &takewhile_type,
4713 &islice_type,
4714 &starmap_type,
4715 &chain_type,
4716 &compress_type,
4717 &filterfalse_type,
4718 &count_type,
4719 &ziplongest_type,
4720 &permutations_type,
4721 &product_type,
4722 &repeat_type,
4723 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004724 &_grouper_type,
4725 &tee_type,
4726 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004727 NULL
4728 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004729
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004730 Py_TYPE(&teedataobject_type) = &PyType_Type;
4731 m = PyModule_Create(&itertoolsmodule);
4732 if (m == NULL)
4733 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004735 for (i=0 ; typelist[i] != NULL ; i++) {
4736 if (PyType_Ready(typelist[i]) < 0)
4737 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004738 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004739 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004740 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004741 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004743 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004744}