blob: 89c0280c9d35c5b1ff484783b3d729fb4d16c8ea [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;
Lisa Roach9718b592018-09-23 17:34:59 -07003478 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003479} accumulateobject;
3480
3481static PyTypeObject accumulate_type;
3482
Tal Einatc4bccd32018-09-12 00:49:13 +03003483/*[clinic input]
3484@classmethod
3485itertools.accumulate.__new__
3486 iterable: object
3487 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003488 *
3489 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003490Return series of accumulated sums (or other binary function results).
3491[clinic start generated code]*/
3492
Raymond Hettinger482ba772010-12-01 22:48:00 +00003493static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003494itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003495 PyObject *binop, PyObject *initial)
3496/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003497{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003498 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003499 accumulateobject *lz;
3500
Raymond Hettinger482ba772010-12-01 22:48:00 +00003501 /* Get iterator. */
3502 it = PyObject_GetIter(iterable);
3503 if (it == NULL)
3504 return NULL;
3505
Raymond Hettinger482ba772010-12-01 22:48:00 +00003506 /* create accumulateobject structure */
3507 lz = (accumulateobject *)type->tp_alloc(type, 0);
3508 if (lz == NULL) {
3509 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003510 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003511 }
3512
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003513 if (binop != Py_None) {
3514 Py_XINCREF(binop);
3515 lz->binop = binop;
3516 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003517 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003518 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003519 Py_XINCREF(initial);
3520 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003521 return (PyObject *)lz;
3522}
3523
3524static void
3525accumulate_dealloc(accumulateobject *lz)
3526{
3527 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003528 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003529 Py_XDECREF(lz->total);
3530 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003531 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003532 Py_TYPE(lz)->tp_free(lz);
3533}
3534
3535static int
3536accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3537{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003538 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003539 Py_VISIT(lz->it);
3540 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003541 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003542 return 0;
3543}
3544
3545static PyObject *
3546accumulate_next(accumulateobject *lz)
3547{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003548 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003549
Lisa Roach9718b592018-09-23 17:34:59 -07003550 if (lz->initial != Py_None) {
3551 lz->total = lz->initial;
3552 Py_INCREF(Py_None);
3553 lz->initial = Py_None;
3554 Py_INCREF(lz->total);
3555 return lz->total;
3556 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003557 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003558 if (val == NULL)
3559 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003560
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003561 if (lz->total == NULL) {
3562 Py_INCREF(val);
3563 lz->total = val;
3564 return lz->total;
3565 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003566
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003567 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003568 newtotal = PyNumber_Add(lz->total, val);
3569 else
3570 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003571 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003572 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003573 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003574
Raymond Hettinger482ba772010-12-01 22:48:00 +00003575 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003576 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003577 return newtotal;
3578}
3579
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003580static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303581accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003582{
Lisa Roach9718b592018-09-23 17:34:59 -07003583 if (lz->initial != Py_None) {
3584 PyObject *it;
3585
3586 assert(lz->total == NULL);
3587 if (PyType_Ready(&chain_type) < 0)
3588 return NULL;
3589 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3590 lz->initial, lz->it);
3591 if (it == NULL)
3592 return NULL;
3593 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3594 it, lz->binop?lz->binop:Py_None, Py_None);
3595 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003596 if (lz->total == Py_None) {
3597 PyObject *it;
3598
3599 if (PyType_Ready(&chain_type) < 0)
3600 return NULL;
3601 if (PyType_Ready(&islice_type) < 0)
3602 return NULL;
3603 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3604 lz->total, lz->it);
3605 if (it == NULL)
3606 return NULL;
3607 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3608 it, lz->binop ? lz->binop : Py_None);
3609 if (it == NULL)
3610 return NULL;
3611 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3612 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003613 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3614 lz->it, lz->binop?lz->binop:Py_None,
3615 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003616}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003617
3618static PyObject *
3619accumulate_setstate(accumulateobject *lz, PyObject *state)
3620{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003621 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003622 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003623 Py_RETURN_NONE;
3624}
3625
3626static PyMethodDef accumulate_methods[] = {
3627 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3628 reduce_doc},
3629 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3630 setstate_doc},
3631 {NULL, NULL} /* sentinel */
3632};
3633
Raymond Hettinger482ba772010-12-01 22:48:00 +00003634static PyTypeObject accumulate_type = {
3635 PyVarObject_HEAD_INIT(NULL, 0)
3636 "itertools.accumulate", /* tp_name */
3637 sizeof(accumulateobject), /* tp_basicsize */
3638 0, /* tp_itemsize */
3639 /* methods */
3640 (destructor)accumulate_dealloc, /* tp_dealloc */
3641 0, /* tp_print */
3642 0, /* tp_getattr */
3643 0, /* tp_setattr */
3644 0, /* tp_reserved */
3645 0, /* tp_repr */
3646 0, /* tp_as_number */
3647 0, /* tp_as_sequence */
3648 0, /* tp_as_mapping */
3649 0, /* tp_hash */
3650 0, /* tp_call */
3651 0, /* tp_str */
3652 PyObject_GenericGetAttr, /* tp_getattro */
3653 0, /* tp_setattro */
3654 0, /* tp_as_buffer */
3655 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3656 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003657 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003658 (traverseproc)accumulate_traverse, /* tp_traverse */
3659 0, /* tp_clear */
3660 0, /* tp_richcompare */
3661 0, /* tp_weaklistoffset */
3662 PyObject_SelfIter, /* tp_iter */
3663 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003664 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003665 0, /* tp_members */
3666 0, /* tp_getset */
3667 0, /* tp_base */
3668 0, /* tp_dict */
3669 0, /* tp_descr_get */
3670 0, /* tp_descr_set */
3671 0, /* tp_dictoffset */
3672 0, /* tp_init */
3673 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003674 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003675 PyObject_GC_Del, /* tp_free */
3676};
3677
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003678
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003679/* compress object ************************************************************/
3680
3681/* Equivalent to:
3682
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003683 def compress(data, selectors):
3684 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3685 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003686*/
3687
3688typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003689 PyObject_HEAD
3690 PyObject *data;
3691 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003692} compressobject;
3693
3694static PyTypeObject compress_type;
3695
Tal Einatc4bccd32018-09-12 00:49:13 +03003696/*[clinic input]
3697@classmethod
3698itertools.compress.__new__
3699 data as seq1: object
3700 selectors as seq2: object
3701Return data elements corresponding to true selector elements.
3702
3703Forms a shorter iterator from selected data elements using the selectors to
3704choose the data elements.
3705[clinic start generated code]*/
3706
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003707static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003708itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3709/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003710{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003711 PyObject *data=NULL, *selectors=NULL;
3712 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003713
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003714 data = PyObject_GetIter(seq1);
3715 if (data == NULL)
3716 goto fail;
3717 selectors = PyObject_GetIter(seq2);
3718 if (selectors == NULL)
3719 goto fail;
3720
3721 /* create compressobject structure */
3722 lz = (compressobject *)type->tp_alloc(type, 0);
3723 if (lz == NULL)
3724 goto fail;
3725 lz->data = data;
3726 lz->selectors = selectors;
3727 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003728
3729fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003730 Py_XDECREF(data);
3731 Py_XDECREF(selectors);
3732 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003733}
3734
3735static void
3736compress_dealloc(compressobject *lz)
3737{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003738 PyObject_GC_UnTrack(lz);
3739 Py_XDECREF(lz->data);
3740 Py_XDECREF(lz->selectors);
3741 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003742}
3743
3744static int
3745compress_traverse(compressobject *lz, visitproc visit, void *arg)
3746{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003747 Py_VISIT(lz->data);
3748 Py_VISIT(lz->selectors);
3749 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003750}
3751
3752static PyObject *
3753compress_next(compressobject *lz)
3754{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003755 PyObject *data = lz->data, *selectors = lz->selectors;
3756 PyObject *datum, *selector;
3757 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3758 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3759 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003760
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003761 while (1) {
3762 /* Steps: get datum, get selector, evaluate selector.
3763 Order is important (to match the pure python version
3764 in terms of which input gets a chance to raise an
3765 exception first).
3766 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003767
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003768 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003769 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003770 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003771
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003772 selector = selectornext(selectors);
3773 if (selector == NULL) {
3774 Py_DECREF(datum);
3775 return NULL;
3776 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003777
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003778 ok = PyObject_IsTrue(selector);
3779 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003780 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003781 return datum;
3782 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003783 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003784 return NULL;
3785 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003786}
3787
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003788static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303789compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003790{
3791 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3792 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003793}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003794
3795static PyMethodDef compress_methods[] = {
3796 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3797 reduce_doc},
3798 {NULL, NULL} /* sentinel */
3799};
3800
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003801static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003802 PyVarObject_HEAD_INIT(NULL, 0)
3803 "itertools.compress", /* tp_name */
3804 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003805 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003806 /* methods */
3807 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003808 0, /* tp_print */
3809 0, /* tp_getattr */
3810 0, /* tp_setattr */
3811 0, /* tp_reserved */
3812 0, /* tp_repr */
3813 0, /* tp_as_number */
3814 0, /* tp_as_sequence */
3815 0, /* tp_as_mapping */
3816 0, /* tp_hash */
3817 0, /* tp_call */
3818 0, /* tp_str */
3819 PyObject_GenericGetAttr, /* tp_getattro */
3820 0, /* tp_setattro */
3821 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003822 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003823 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003824 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003825 (traverseproc)compress_traverse, /* tp_traverse */
3826 0, /* tp_clear */
3827 0, /* tp_richcompare */
3828 0, /* tp_weaklistoffset */
3829 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003830 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003831 compress_methods, /* tp_methods */
3832 0, /* tp_members */
3833 0, /* tp_getset */
3834 0, /* tp_base */
3835 0, /* tp_dict */
3836 0, /* tp_descr_get */
3837 0, /* tp_descr_set */
3838 0, /* tp_dictoffset */
3839 0, /* tp_init */
3840 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003841 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003842 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003843};
3844
3845
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003846/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003847
3848typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003849 PyObject_HEAD
3850 PyObject *func;
3851 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003852} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003853
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003854static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003855
Tal Einatc4bccd32018-09-12 00:49:13 +03003856/*[clinic input]
3857@classmethod
3858itertools.filterfalse.__new__
3859 function as func: object
3860 iterable as seq: object
3861 /
3862Return those items of iterable for which function(item) is false.
3863
3864If function is None, return the items that are false.
3865[clinic start generated code]*/
3866
Raymond Hettinger60eca932003-02-09 06:40:58 +00003867static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003868itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3869/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003870{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003871 PyObject *it;
3872 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003873
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003874 /* Get iterator. */
3875 it = PyObject_GetIter(seq);
3876 if (it == NULL)
3877 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003879 /* create filterfalseobject structure */
3880 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3881 if (lz == NULL) {
3882 Py_DECREF(it);
3883 return NULL;
3884 }
3885 Py_INCREF(func);
3886 lz->func = func;
3887 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003889 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003890}
3891
3892static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003893filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003894{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003895 PyObject_GC_UnTrack(lz);
3896 Py_XDECREF(lz->func);
3897 Py_XDECREF(lz->it);
3898 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003899}
3900
3901static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003902filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003903{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003904 Py_VISIT(lz->it);
3905 Py_VISIT(lz->func);
3906 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003907}
3908
3909static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003910filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003911{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003912 PyObject *item;
3913 PyObject *it = lz->it;
3914 long ok;
3915 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003916
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003917 iternext = *Py_TYPE(it)->tp_iternext;
3918 for (;;) {
3919 item = iternext(it);
3920 if (item == NULL)
3921 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003922
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003923 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3924 ok = PyObject_IsTrue(item);
3925 } else {
3926 PyObject *good;
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01003927 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003928 if (good == NULL) {
3929 Py_DECREF(item);
3930 return NULL;
3931 }
3932 ok = PyObject_IsTrue(good);
3933 Py_DECREF(good);
3934 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003935 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003936 return item;
3937 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003938 if (ok < 0)
3939 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003940 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003941}
3942
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003943static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303944filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003945{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003946 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3947}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003948
3949static PyMethodDef filterfalse_methods[] = {
3950 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3951 reduce_doc},
3952 {NULL, NULL} /* sentinel */
3953};
3954
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003955static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003956 PyVarObject_HEAD_INIT(NULL, 0)
3957 "itertools.filterfalse", /* tp_name */
3958 sizeof(filterfalseobject), /* tp_basicsize */
3959 0, /* tp_itemsize */
3960 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003961 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003962 0, /* tp_print */
3963 0, /* tp_getattr */
3964 0, /* tp_setattr */
3965 0, /* tp_reserved */
3966 0, /* tp_repr */
3967 0, /* tp_as_number */
3968 0, /* tp_as_sequence */
3969 0, /* tp_as_mapping */
3970 0, /* tp_hash */
3971 0, /* tp_call */
3972 0, /* tp_str */
3973 PyObject_GenericGetAttr, /* tp_getattro */
3974 0, /* tp_setattro */
3975 0, /* tp_as_buffer */
3976 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3977 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003978 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003979 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003980 0, /* tp_clear */
3981 0, /* tp_richcompare */
3982 0, /* tp_weaklistoffset */
3983 PyObject_SelfIter, /* tp_iter */
3984 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003985 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003986 0, /* tp_members */
3987 0, /* tp_getset */
3988 0, /* tp_base */
3989 0, /* tp_dict */
3990 0, /* tp_descr_get */
3991 0, /* tp_descr_set */
3992 0, /* tp_dictoffset */
3993 0, /* tp_init */
3994 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003995 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003996 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003997};
3998
3999
4000/* count object ************************************************************/
4001
4002typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004003 PyObject_HEAD
4004 Py_ssize_t cnt;
4005 PyObject *long_cnt;
4006 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004007} countobject;
4008
Raymond Hettinger30729212009-02-12 06:28:27 +00004009/* Counting logic and invariants:
4010
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004011fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00004012
Serhiy Storchaka483405b2015-02-17 10:14:30 +02004013 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004014 Advances with: cnt += 1
4015 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00004016
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004017slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004018
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004019 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4020 All counting is done with python objects (no overflows or underflows).
4021 Advances with: long_cnt += long_step
4022 Step may be zero -- effectively a slow version of repeat(cnt).
4023 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004024*/
4025
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004026static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004027
Tal Einatc4bccd32018-09-12 00:49:13 +03004028/*[clinic input]
4029@classmethod
4030itertools.count.__new__
4031 start as long_cnt: object(c_default="NULL") = 0
4032 step as long_step: object(c_default="NULL") = 1
4033Return a count object whose .__next__() method returns consecutive values.
4034
4035Equivalent to:
4036 def count(firstval=0, step=1):
4037 x = firstval
4038 while 1:
4039 yield x
4040 x += step
4041[clinic start generated code]*/
4042
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004043static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004044itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4045 PyObject *long_step)
4046/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004047{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004048 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004049 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004050 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004051 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004053 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4054 (long_step != NULL && !PyNumber_Check(long_step))) {
4055 PyErr_SetString(PyExc_TypeError, "a number is required");
4056 return NULL;
4057 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004058
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004059 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4060 (long_step == NULL || PyLong_Check(long_step));
4061
4062 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004063 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004064 if (fast_mode) {
4065 assert(PyLong_Check(long_cnt));
4066 cnt = PyLong_AsSsize_t(long_cnt);
4067 if (cnt == -1 && PyErr_Occurred()) {
4068 PyErr_Clear();
4069 fast_mode = 0;
4070 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004071 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004072 } else {
4073 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004074 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004075 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004076 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004077
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004078 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004079 if (long_step == NULL)
4080 long_step = _PyLong_One;
4081 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004082
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004083 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004084
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004085 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004086 if (fast_mode) {
4087 assert(PyLong_Check(long_step));
4088 step = PyLong_AsLong(long_step);
4089 if (step != 1) {
4090 fast_mode = 0;
4091 if (step == -1 && PyErr_Occurred())
4092 PyErr_Clear();
4093 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004094 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004095
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004096 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004097 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004098 else
4099 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004100
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004101 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4102 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4103 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004104 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004105
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004106 /* create countobject structure */
4107 lz = (countobject *)type->tp_alloc(type, 0);
4108 if (lz == NULL) {
4109 Py_XDECREF(long_cnt);
4110 return NULL;
4111 }
4112 lz->cnt = cnt;
4113 lz->long_cnt = long_cnt;
4114 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004115
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004116 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004117}
4118
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004119static void
4120count_dealloc(countobject *lz)
4121{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004122 PyObject_GC_UnTrack(lz);
4123 Py_XDECREF(lz->long_cnt);
4124 Py_XDECREF(lz->long_step);
4125 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004126}
4127
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004128static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004129count_traverse(countobject *lz, visitproc visit, void *arg)
4130{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004131 Py_VISIT(lz->long_cnt);
4132 Py_VISIT(lz->long_step);
4133 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004134}
4135
4136static PyObject *
4137count_nextlong(countobject *lz)
4138{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004139 PyObject *long_cnt;
4140 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004141
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004142 long_cnt = lz->long_cnt;
4143 if (long_cnt == NULL) {
4144 /* Switch to slow_mode */
4145 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4146 if (long_cnt == NULL)
4147 return NULL;
4148 }
4149 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004150
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004151 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4152 if (stepped_up == NULL)
4153 return NULL;
4154 lz->long_cnt = stepped_up;
4155 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004156}
4157
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004158static PyObject *
4159count_next(countobject *lz)
4160{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004161 if (lz->cnt == PY_SSIZE_T_MAX)
4162 return count_nextlong(lz);
4163 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004164}
4165
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004166static PyObject *
4167count_repr(countobject *lz)
4168{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004169 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004170 return PyUnicode_FromFormat("%s(%zd)",
4171 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004172
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004173 if (PyLong_Check(lz->long_step)) {
4174 long step = PyLong_AsLong(lz->long_step);
4175 if (step == -1 && PyErr_Occurred()) {
4176 PyErr_Clear();
4177 }
4178 if (step == 1) {
4179 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004180 return PyUnicode_FromFormat("%s(%R)",
4181 _PyType_Name(Py_TYPE(lz)),
4182 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004183 }
4184 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004185 return PyUnicode_FromFormat("%s(%R, %R)",
4186 _PyType_Name(Py_TYPE(lz)),
4187 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004188}
4189
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004190static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304191count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004192{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004193 if (lz->cnt == PY_SSIZE_T_MAX)
4194 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4195 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004196}
4197
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004198static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004199 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004200 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004201 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004202};
4203
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004204static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004205 PyVarObject_HEAD_INIT(NULL, 0)
4206 "itertools.count", /* tp_name */
4207 sizeof(countobject), /* tp_basicsize */
4208 0, /* tp_itemsize */
4209 /* methods */
4210 (destructor)count_dealloc, /* tp_dealloc */
4211 0, /* tp_print */
4212 0, /* tp_getattr */
4213 0, /* tp_setattr */
4214 0, /* tp_reserved */
4215 (reprfunc)count_repr, /* tp_repr */
4216 0, /* tp_as_number */
4217 0, /* tp_as_sequence */
4218 0, /* tp_as_mapping */
4219 0, /* tp_hash */
4220 0, /* tp_call */
4221 0, /* tp_str */
4222 PyObject_GenericGetAttr, /* tp_getattro */
4223 0, /* tp_setattro */
4224 0, /* tp_as_buffer */
4225 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004226 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004227 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004228 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004229 0, /* tp_clear */
4230 0, /* tp_richcompare */
4231 0, /* tp_weaklistoffset */
4232 PyObject_SelfIter, /* tp_iter */
4233 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004234 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004235 0, /* tp_members */
4236 0, /* tp_getset */
4237 0, /* tp_base */
4238 0, /* tp_dict */
4239 0, /* tp_descr_get */
4240 0, /* tp_descr_set */
4241 0, /* tp_dictoffset */
4242 0, /* tp_init */
4243 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004244 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004245 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004246};
4247
4248
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004249/* repeat object ************************************************************/
4250
4251typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004252 PyObject_HEAD
4253 PyObject *element;
4254 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004255} repeatobject;
4256
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004257static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004258
4259static PyObject *
4260repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4261{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004262 repeatobject *ro;
4263 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004264 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004265 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004267 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4268 &element, &cnt))
4269 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004270
Raymond Hettinger97d35552014-06-24 21:36:58 -07004271 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004272 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004273 /* Does user supply times argument? */
4274 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004275 cnt = 0;
4276
4277 ro = (repeatobject *)type->tp_alloc(type, 0);
4278 if (ro == NULL)
4279 return NULL;
4280 Py_INCREF(element);
4281 ro->element = element;
4282 ro->cnt = cnt;
4283 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004284}
4285
4286static void
4287repeat_dealloc(repeatobject *ro)
4288{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004289 PyObject_GC_UnTrack(ro);
4290 Py_XDECREF(ro->element);
4291 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004292}
4293
4294static int
4295repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4296{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004297 Py_VISIT(ro->element);
4298 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004299}
4300
4301static PyObject *
4302repeat_next(repeatobject *ro)
4303{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004304 if (ro->cnt == 0)
4305 return NULL;
4306 if (ro->cnt > 0)
4307 ro->cnt--;
4308 Py_INCREF(ro->element);
4309 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004310}
4311
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004312static PyObject *
4313repeat_repr(repeatobject *ro)
4314{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004315 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004316 return PyUnicode_FromFormat("%s(%R)",
4317 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004318 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004319 return PyUnicode_FromFormat("%s(%R, %zd)",
4320 _PyType_Name(Py_TYPE(ro)), ro->element,
4321 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004322}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004323
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004324static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304325repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004326{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004327 if (ro->cnt == -1) {
4328 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4329 return NULL;
4330 }
4331 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004332}
4333
Armin Rigof5b3e362006-02-11 21:32:43 +00004334PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004335
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004336static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304337repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004338{
4339 /* unpickle this so that a new repeat iterator is constructed with an
4340 * object, then call __setstate__ on it to set cnt
4341 */
4342 if (ro->cnt >= 0)
4343 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4344 else
4345 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4346}
4347
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004348static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004349 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004350 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004351 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004352};
4353
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004354PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004355"repeat(object [,times]) -> create an iterator which returns the object\n\
4356for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004357endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004358
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004359static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004360 PyVarObject_HEAD_INIT(NULL, 0)
4361 "itertools.repeat", /* tp_name */
4362 sizeof(repeatobject), /* tp_basicsize */
4363 0, /* tp_itemsize */
4364 /* methods */
4365 (destructor)repeat_dealloc, /* tp_dealloc */
4366 0, /* tp_print */
4367 0, /* tp_getattr */
4368 0, /* tp_setattr */
4369 0, /* tp_reserved */
4370 (reprfunc)repeat_repr, /* tp_repr */
4371 0, /* tp_as_number */
4372 0, /* tp_as_sequence */
4373 0, /* tp_as_mapping */
4374 0, /* tp_hash */
4375 0, /* tp_call */
4376 0, /* tp_str */
4377 PyObject_GenericGetAttr, /* tp_getattro */
4378 0, /* tp_setattro */
4379 0, /* tp_as_buffer */
4380 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4381 Py_TPFLAGS_BASETYPE, /* tp_flags */
4382 repeat_doc, /* tp_doc */
4383 (traverseproc)repeat_traverse, /* tp_traverse */
4384 0, /* tp_clear */
4385 0, /* tp_richcompare */
4386 0, /* tp_weaklistoffset */
4387 PyObject_SelfIter, /* tp_iter */
4388 (iternextfunc)repeat_next, /* tp_iternext */
4389 repeat_methods, /* tp_methods */
4390 0, /* tp_members */
4391 0, /* tp_getset */
4392 0, /* tp_base */
4393 0, /* tp_dict */
4394 0, /* tp_descr_get */
4395 0, /* tp_descr_set */
4396 0, /* tp_dictoffset */
4397 0, /* tp_init */
4398 0, /* tp_alloc */
4399 repeat_new, /* tp_new */
4400 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004401};
4402
Tal Einatc4bccd32018-09-12 00:49:13 +03004403
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004404/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004405
4406typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004407 PyObject_HEAD
4408 Py_ssize_t tuplesize;
4409 Py_ssize_t numactive;
4410 PyObject *ittuple; /* tuple of iterators */
4411 PyObject *result;
4412 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004413} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004414
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004415static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004416
4417static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004418zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004419{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004420 ziplongestobject *lz;
4421 Py_ssize_t i;
4422 PyObject *ittuple; /* tuple of iterators */
4423 PyObject *result;
4424 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004425 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004426
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004427 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004428 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004429 if (fillvalue == NULL || PyDict_GET_SIZE(kwds) > 1) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004430 PyErr_SetString(PyExc_TypeError,
4431 "zip_longest() got an unexpected keyword argument");
4432 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004433 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004434 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004435
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004436 /* args must be a tuple */
4437 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004438 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004439
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004440 /* obtain iterators */
4441 ittuple = PyTuple_New(tuplesize);
4442 if (ittuple == NULL)
4443 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004444 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004445 PyObject *item = PyTuple_GET_ITEM(args, i);
4446 PyObject *it = PyObject_GetIter(item);
4447 if (it == NULL) {
4448 if (PyErr_ExceptionMatches(PyExc_TypeError))
4449 PyErr_Format(PyExc_TypeError,
4450 "zip_longest argument #%zd must support iteration",
4451 i+1);
4452 Py_DECREF(ittuple);
4453 return NULL;
4454 }
4455 PyTuple_SET_ITEM(ittuple, i, it);
4456 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004457
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004458 /* create a result holder */
4459 result = PyTuple_New(tuplesize);
4460 if (result == NULL) {
4461 Py_DECREF(ittuple);
4462 return NULL;
4463 }
4464 for (i=0 ; i < tuplesize ; i++) {
4465 Py_INCREF(Py_None);
4466 PyTuple_SET_ITEM(result, i, Py_None);
4467 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004468
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004469 /* create ziplongestobject structure */
4470 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4471 if (lz == NULL) {
4472 Py_DECREF(ittuple);
4473 Py_DECREF(result);
4474 return NULL;
4475 }
4476 lz->ittuple = ittuple;
4477 lz->tuplesize = tuplesize;
4478 lz->numactive = tuplesize;
4479 lz->result = result;
4480 Py_INCREF(fillvalue);
4481 lz->fillvalue = fillvalue;
4482 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004483}
4484
4485static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004486zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004487{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004488 PyObject_GC_UnTrack(lz);
4489 Py_XDECREF(lz->ittuple);
4490 Py_XDECREF(lz->result);
4491 Py_XDECREF(lz->fillvalue);
4492 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004493}
4494
4495static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004496zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004497{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004498 Py_VISIT(lz->ittuple);
4499 Py_VISIT(lz->result);
4500 Py_VISIT(lz->fillvalue);
4501 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004502}
4503
4504static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004505zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004507 Py_ssize_t i;
4508 Py_ssize_t tuplesize = lz->tuplesize;
4509 PyObject *result = lz->result;
4510 PyObject *it;
4511 PyObject *item;
4512 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004513
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004514 if (tuplesize == 0)
4515 return NULL;
4516 if (lz->numactive == 0)
4517 return NULL;
4518 if (Py_REFCNT(result) == 1) {
4519 Py_INCREF(result);
4520 for (i=0 ; i < tuplesize ; i++) {
4521 it = PyTuple_GET_ITEM(lz->ittuple, i);
4522 if (it == NULL) {
4523 Py_INCREF(lz->fillvalue);
4524 item = lz->fillvalue;
4525 } else {
4526 item = PyIter_Next(it);
4527 if (item == NULL) {
4528 lz->numactive -= 1;
4529 if (lz->numactive == 0 || PyErr_Occurred()) {
4530 lz->numactive = 0;
4531 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004532 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004533 } else {
4534 Py_INCREF(lz->fillvalue);
4535 item = lz->fillvalue;
4536 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4537 Py_DECREF(it);
4538 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004539 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004540 }
4541 olditem = PyTuple_GET_ITEM(result, i);
4542 PyTuple_SET_ITEM(result, i, item);
4543 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004544 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004545 } else {
4546 result = PyTuple_New(tuplesize);
4547 if (result == NULL)
4548 return NULL;
4549 for (i=0 ; i < tuplesize ; i++) {
4550 it = PyTuple_GET_ITEM(lz->ittuple, i);
4551 if (it == NULL) {
4552 Py_INCREF(lz->fillvalue);
4553 item = lz->fillvalue;
4554 } else {
4555 item = PyIter_Next(it);
4556 if (item == NULL) {
4557 lz->numactive -= 1;
4558 if (lz->numactive == 0 || PyErr_Occurred()) {
4559 lz->numactive = 0;
4560 Py_DECREF(result);
4561 return NULL;
4562 } else {
4563 Py_INCREF(lz->fillvalue);
4564 item = lz->fillvalue;
4565 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4566 Py_DECREF(it);
4567 }
4568 }
4569 }
4570 PyTuple_SET_ITEM(result, i, item);
4571 }
4572 }
4573 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004574}
4575
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004576static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304577zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004578{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004579
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004580 /* Create a new tuple with empty sequences where appropriate to pickle.
4581 * Then use setstate to set the fillvalue
4582 */
4583 int i;
4584 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004585
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004586 if (args == NULL)
4587 return NULL;
4588 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4589 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4590 if (elem == NULL) {
4591 elem = PyTuple_New(0);
4592 if (elem == NULL) {
4593 Py_DECREF(args);
4594 return NULL;
4595 }
4596 } else
4597 Py_INCREF(elem);
4598 PyTuple_SET_ITEM(args, i, elem);
4599 }
4600 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4601}
4602
4603static PyObject *
4604zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4605{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004606 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004607 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004608 Py_RETURN_NONE;
4609}
4610
4611static PyMethodDef zip_longest_methods[] = {
4612 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4613 reduce_doc},
4614 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4615 setstate_doc},
4616 {NULL, NULL} /* sentinel */
4617};
4618
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004619PyDoc_STRVAR(zip_longest_doc,
4620"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004621\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004622Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004623the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004624method continues until the longest iterable in the argument sequence\n\
4625is exhausted and then it raises StopIteration. When the shorter iterables\n\
4626are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4627defaults to None or can be specified by a keyword argument.\n\
4628");
4629
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004630static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004631 PyVarObject_HEAD_INIT(NULL, 0)
4632 "itertools.zip_longest", /* tp_name */
4633 sizeof(ziplongestobject), /* tp_basicsize */
4634 0, /* tp_itemsize */
4635 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004636 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004637 0, /* tp_print */
4638 0, /* tp_getattr */
4639 0, /* tp_setattr */
4640 0, /* tp_reserved */
4641 0, /* tp_repr */
4642 0, /* tp_as_number */
4643 0, /* tp_as_sequence */
4644 0, /* tp_as_mapping */
4645 0, /* tp_hash */
4646 0, /* tp_call */
4647 0, /* tp_str */
4648 PyObject_GenericGetAttr, /* tp_getattro */
4649 0, /* tp_setattro */
4650 0, /* tp_as_buffer */
4651 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4652 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004653 zip_longest_doc, /* tp_doc */
4654 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004655 0, /* tp_clear */
4656 0, /* tp_richcompare */
4657 0, /* tp_weaklistoffset */
4658 PyObject_SelfIter, /* tp_iter */
4659 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004660 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004661 0, /* tp_members */
4662 0, /* tp_getset */
4663 0, /* tp_base */
4664 0, /* tp_dict */
4665 0, /* tp_descr_get */
4666 0, /* tp_descr_set */
4667 0, /* tp_dictoffset */
4668 0, /* tp_init */
4669 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004670 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004671 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004672};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004673
Tal Einatc4bccd32018-09-12 00:49:13 +03004674
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004675/* module level code ********************************************************/
4676
4677PyDoc_STRVAR(module_doc,
4678"Functional tools for creating and using iterators.\n\
4679\n\
4680Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004681count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004682cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004683repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004684\n\
4685Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004686accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004687chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4688chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004689compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4690dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4691groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004692filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004693islice(seq, [start,] stop [, step]) --> elements from\n\
4694 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004695starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004696tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004697takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004698zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004699\n\
4700Combinatoric generators:\n\
4701product(p, q, ... [repeat=1]) --> cartesian product\n\
4702permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004703combinations(p, r)\n\
4704combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004705");
4706
4707
Raymond Hettingerad983e72003-11-12 14:32:26 +00004708static PyMethodDef module_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03004709 ITERTOOLS_TEE_METHODDEF
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004710 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004711};
4712
Martin v. Löwis1a214512008-06-11 05:26:20 +00004713
4714static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004715 PyModuleDef_HEAD_INIT,
4716 "itertools",
4717 module_doc,
4718 -1,
4719 module_methods,
4720 NULL,
4721 NULL,
4722 NULL,
4723 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004724};
4725
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004726PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004727PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004729 int i;
4730 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004731 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004732 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004733 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004734 &combinations_type,
4735 &cwr_type,
4736 &cycle_type,
4737 &dropwhile_type,
4738 &takewhile_type,
4739 &islice_type,
4740 &starmap_type,
4741 &chain_type,
4742 &compress_type,
4743 &filterfalse_type,
4744 &count_type,
4745 &ziplongest_type,
4746 &permutations_type,
4747 &product_type,
4748 &repeat_type,
4749 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004750 &_grouper_type,
4751 &tee_type,
4752 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004753 NULL
4754 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004756 Py_TYPE(&teedataobject_type) = &PyType_Type;
4757 m = PyModule_Create(&itertoolsmodule);
4758 if (m == NULL)
4759 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004760
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004761 for (i=0 ; typelist[i] != NULL ; i++) {
4762 if (PyType_Ready(typelist[i]) < 0)
4763 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004764 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004765 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004766 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004767 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004768
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004769 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004770}