blob: 581ae2c67e853c5688ae2322001fc65236de85f4 [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 Storchakabb86bf42018-12-11 08:28:18 +0200371 _Py_IDENTIFIER(iter);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300372 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200373 return Py_BuildValue("N(())", _PyEval_GetBuiltinId(&PyId_iter));
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300374 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700375 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000376}
377
378static PyMethodDef _grouper_methods[] = {
379 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
380 reduce_doc},
381 {NULL, NULL} /* sentinel */
382};
383
384
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000385static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000386 PyVarObject_HEAD_INIT(NULL, 0)
387 "itertools._grouper", /* tp_name */
388 sizeof(_grouperobject), /* tp_basicsize */
389 0, /* tp_itemsize */
390 /* methods */
391 (destructor)_grouper_dealloc, /* tp_dealloc */
392 0, /* tp_print */
393 0, /* tp_getattr */
394 0, /* tp_setattr */
395 0, /* tp_reserved */
396 0, /* tp_repr */
397 0, /* tp_as_number */
398 0, /* tp_as_sequence */
399 0, /* tp_as_mapping */
400 0, /* tp_hash */
401 0, /* tp_call */
402 0, /* tp_str */
403 PyObject_GenericGetAttr, /* tp_getattro */
404 0, /* tp_setattro */
405 0, /* tp_as_buffer */
406 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
407 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700408 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000409 0, /* tp_clear */
410 0, /* tp_richcompare */
411 0, /* tp_weaklistoffset */
412 PyObject_SelfIter, /* tp_iter */
413 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000414 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000415 0, /* tp_members */
416 0, /* tp_getset */
417 0, /* tp_base */
418 0, /* tp_dict */
419 0, /* tp_descr_get */
420 0, /* tp_descr_set */
421 0, /* tp_dictoffset */
422 0, /* tp_init */
423 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300424 itertools__grouper, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000425 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000426};
427
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000428
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700429/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000430
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431/* The teedataobject pre-allocates space for LINKCELLS number of objects.
432 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000433 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000434 the other structure members including PyHEAD overhead. The larger the
435 value, the less memory overhead per object and the less time spent
436 allocating/deallocating new links. The smaller the number, the less
437 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000438*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000439#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000440
441typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000442 PyObject_HEAD
443 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700444 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000445 PyObject *nextlink;
446 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000447} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000448
449typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000450 PyObject_HEAD
451 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700452 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000453 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000454} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000455
Raymond Hettingerad983e72003-11-12 14:32:26 +0000456static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000457
458static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000459teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000461 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
464 if (tdo == NULL)
465 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000467 tdo->numread = 0;
468 tdo->nextlink = NULL;
469 Py_INCREF(it);
470 tdo->it = it;
471 PyObject_GC_Track(tdo);
472 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000473}
474
475static PyObject *
476teedataobject_jumplink(teedataobject *tdo)
477{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000478 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000479 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000480 Py_XINCREF(tdo->nextlink);
481 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000482}
483
484static PyObject *
485teedataobject_getitem(teedataobject *tdo, int i)
486{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000487 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000489 assert(i < LINKCELLS);
490 if (i < tdo->numread)
491 value = tdo->values[i];
492 else {
493 /* this is the lead iterator, so fetch more data */
494 assert(i == tdo->numread);
495 value = PyIter_Next(tdo->it);
496 if (value == NULL)
497 return NULL;
498 tdo->numread++;
499 tdo->values[i] = value;
500 }
501 Py_INCREF(value);
502 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000503}
504
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000505static int
506teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
507{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000508 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700509
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000510 Py_VISIT(tdo->it);
511 for (i = 0; i < tdo->numread; i++)
512 Py_VISIT(tdo->values[i]);
513 Py_VISIT(tdo->nextlink);
514 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000515}
516
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200517static void
518teedataobject_safe_decref(PyObject *obj)
519{
520 while (obj && Py_TYPE(obj) == &teedataobject_type &&
521 Py_REFCNT(obj) == 1) {
522 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
523 ((teedataobject *)obj)->nextlink = NULL;
524 Py_DECREF(obj);
525 obj = nextlink;
526 }
527 Py_XDECREF(obj);
528}
529
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000530static int
531teedataobject_clear(teedataobject *tdo)
532{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000533 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200534 PyObject *tmp;
535
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000536 Py_CLEAR(tdo->it);
537 for (i=0 ; i<tdo->numread ; i++)
538 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200539 tmp = tdo->nextlink;
540 tdo->nextlink = NULL;
541 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000542 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000543}
544
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000545static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000546teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000547{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000548 PyObject_GC_UnTrack(tdo);
549 teedataobject_clear(tdo);
550 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000551}
552
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000553static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530554teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000555{
556 int i;
557 /* create a temporary list of already iterated values */
558 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700559
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000560 if (!values)
561 return NULL;
562 for (i=0 ; i<tdo->numread ; i++) {
563 Py_INCREF(tdo->values[i]);
564 PyList_SET_ITEM(values, i, tdo->values[i]);
565 }
566 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
567 values,
568 tdo->nextlink ? tdo->nextlink : Py_None);
569}
570
Tal Einatc4bccd32018-09-12 00:49:13 +0300571/*[clinic input]
572@classmethod
573itertools.teedataobject.__new__
574 iterable as it: object
575 values: object(subclass_of='&PyList_Type')
576 next: object
577 /
578Data container common to multiple tee objects.
579[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000580
581static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300582itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
583 PyObject *values, PyObject *next)
584/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000585{
586 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000587 Py_ssize_t i, len;
588
589 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000590
591 tdo = (teedataobject *)teedataobject_newinternal(it);
592 if (!tdo)
593 return NULL;
594
595 len = PyList_GET_SIZE(values);
596 if (len > LINKCELLS)
597 goto err;
598 for (i=0; i<len; i++) {
599 tdo->values[i] = PyList_GET_ITEM(values, i);
600 Py_INCREF(tdo->values[i]);
601 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200602 /* len <= LINKCELLS < INT_MAX */
603 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000604
605 if (len == LINKCELLS) {
606 if (next != Py_None) {
607 if (Py_TYPE(next) != &teedataobject_type)
608 goto err;
609 assert(tdo->nextlink == NULL);
610 Py_INCREF(next);
611 tdo->nextlink = next;
612 }
613 } else {
614 if (next != Py_None)
615 goto err; /* shouldn't have a next if we are not full */
616 }
617 return (PyObject*)tdo;
618
619err:
620 Py_XDECREF(tdo);
621 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
622 return NULL;
623}
624
625static PyMethodDef teedataobject_methods[] = {
626 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
627 reduce_doc},
628 {NULL, NULL} /* sentinel */
629};
630
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700632 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000633 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000634 sizeof(teedataobject), /* tp_basicsize */
635 0, /* tp_itemsize */
636 /* methods */
637 (destructor)teedataobject_dealloc, /* tp_dealloc */
638 0, /* tp_print */
639 0, /* tp_getattr */
640 0, /* tp_setattr */
641 0, /* tp_reserved */
642 0, /* tp_repr */
643 0, /* tp_as_number */
644 0, /* tp_as_sequence */
645 0, /* tp_as_mapping */
646 0, /* tp_hash */
647 0, /* tp_call */
648 0, /* tp_str */
649 PyObject_GenericGetAttr, /* tp_getattro */
650 0, /* tp_setattro */
651 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700652 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300653 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000654 (traverseproc)teedataobject_traverse, /* tp_traverse */
655 (inquiry)teedataobject_clear, /* tp_clear */
656 0, /* tp_richcompare */
657 0, /* tp_weaklistoffset */
658 0, /* tp_iter */
659 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000660 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000661 0, /* tp_members */
662 0, /* tp_getset */
663 0, /* tp_base */
664 0, /* tp_dict */
665 0, /* tp_descr_get */
666 0, /* tp_descr_set */
667 0, /* tp_dictoffset */
668 0, /* tp_init */
669 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300670 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000671 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000672};
673
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000674
675static PyTypeObject tee_type;
676
677static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000678tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000679{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000680 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000681
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000682 if (to->index >= LINKCELLS) {
683 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200684 if (link == NULL)
685 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300686 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 to->index = 0;
688 }
689 value = teedataobject_getitem(to->dataobj, to->index);
690 if (value == NULL)
691 return NULL;
692 to->index++;
693 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000694}
695
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000696static int
697tee_traverse(teeobject *to, visitproc visit, void *arg)
698{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000699 Py_VISIT((PyObject *)to->dataobj);
700 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000701}
702
Raymond Hettingerad983e72003-11-12 14:32:26 +0000703static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530704tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000705{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000706 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000707
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000708 newto = PyObject_GC_New(teeobject, &tee_type);
709 if (newto == NULL)
710 return NULL;
711 Py_INCREF(to->dataobj);
712 newto->dataobj = to->dataobj;
713 newto->index = to->index;
714 newto->weakreflist = NULL;
715 PyObject_GC_Track(newto);
716 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000717}
718
719PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
720
721static PyObject *
722tee_fromiterable(PyObject *iterable)
723{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000724 teeobject *to;
725 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000726
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000727 it = PyObject_GetIter(iterable);
728 if (it == NULL)
729 return NULL;
730 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530731 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000732 goto done;
733 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000735 to = PyObject_GC_New(teeobject, &tee_type);
736 if (to == NULL)
737 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000738 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000739 if (!to->dataobj) {
740 PyObject_GC_Del(to);
741 to = NULL;
742 goto done;
743 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000745 to->index = 0;
746 to->weakreflist = NULL;
747 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000748done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000749 Py_XDECREF(it);
750 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000751}
752
Tal Einatc4bccd32018-09-12 00:49:13 +0300753/*[clinic input]
754@classmethod
755itertools._tee.__new__
756 iterable: object
757 /
758Iterator wrapped to make it copyable.
759[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760
Tal Einatc4bccd32018-09-12 00:49:13 +0300761static PyObject *
762itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
763/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
764{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000765 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000766}
767
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000768static int
769tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000770{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000771 if (to->weakreflist != NULL)
772 PyObject_ClearWeakRefs((PyObject *) to);
773 Py_CLEAR(to->dataobj);
774 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000775}
776
777static void
778tee_dealloc(teeobject *to)
779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000780 PyObject_GC_UnTrack(to);
781 tee_clear(to);
782 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000783}
784
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000785static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530786tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000787{
788 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
789}
790
791static PyObject *
792tee_setstate(teeobject *to, PyObject *state)
793{
794 teedataobject *tdo;
795 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300796 if (!PyTuple_Check(state)) {
797 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000798 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300799 }
800 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
801 return NULL;
802 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000803 if (index < 0 || index > LINKCELLS) {
804 PyErr_SetString(PyExc_ValueError, "Index out of range");
805 return NULL;
806 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200807 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300808 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000809 to->index = index;
810 Py_RETURN_NONE;
811}
812
Raymond Hettingerad983e72003-11-12 14:32:26 +0000813static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700814 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
815 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
816 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000817 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000818};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000819
820static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000821 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000822 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000823 sizeof(teeobject), /* tp_basicsize */
824 0, /* tp_itemsize */
825 /* methods */
826 (destructor)tee_dealloc, /* tp_dealloc */
827 0, /* tp_print */
828 0, /* tp_getattr */
829 0, /* tp_setattr */
830 0, /* tp_reserved */
831 0, /* tp_repr */
832 0, /* tp_as_number */
833 0, /* tp_as_sequence */
834 0, /* tp_as_mapping */
835 0, /* tp_hash */
836 0, /* tp_call */
837 0, /* tp_str */
838 0, /* tp_getattro */
839 0, /* tp_setattro */
840 0, /* tp_as_buffer */
841 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300842 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000843 (traverseproc)tee_traverse, /* tp_traverse */
844 (inquiry)tee_clear, /* tp_clear */
845 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700846 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 PyObject_SelfIter, /* tp_iter */
848 (iternextfunc)tee_next, /* tp_iternext */
849 tee_methods, /* tp_methods */
850 0, /* tp_members */
851 0, /* tp_getset */
852 0, /* tp_base */
853 0, /* tp_dict */
854 0, /* tp_descr_get */
855 0, /* tp_descr_set */
856 0, /* tp_dictoffset */
857 0, /* tp_init */
858 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300859 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000860 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000861};
862
Tal Einatc4bccd32018-09-12 00:49:13 +0300863/*[clinic input]
864itertools.tee
865 iterable: object
866 n: Py_ssize_t = 2
867 /
868Returns a tuple of n independent iterators.
869[clinic start generated code]*/
870
Raymond Hettingerad983e72003-11-12 14:32:26 +0000871static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300872itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
873/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000874{
Tal Einatc4bccd32018-09-12 00:49:13 +0300875 Py_ssize_t i;
876 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200877 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000879 if (n < 0) {
880 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
881 return NULL;
882 }
883 result = PyTuple_New(n);
884 if (result == NULL)
885 return NULL;
886 if (n == 0)
887 return result;
888 it = PyObject_GetIter(iterable);
889 if (it == NULL) {
890 Py_DECREF(result);
891 return NULL;
892 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200893
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200894 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200895 Py_DECREF(it);
896 Py_DECREF(result);
897 return NULL;
898 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200899 if (copyfunc != NULL) {
900 copyable = it;
901 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200902 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000903 copyable = tee_fromiterable(it);
904 Py_DECREF(it);
905 if (copyable == NULL) {
906 Py_DECREF(result);
907 return NULL;
908 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200909 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
910 if (copyfunc == NULL) {
911 Py_DECREF(copyable);
912 Py_DECREF(result);
913 return NULL;
914 }
915 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200916
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200917 PyTuple_SET_ITEM(result, 0, copyable);
918 for (i = 1; i < n; i++) {
919 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000920 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200921 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000922 Py_DECREF(result);
923 return NULL;
924 }
925 PyTuple_SET_ITEM(result, i, copyable);
926 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200927 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000928 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000929}
930
Raymond Hettingerad983e72003-11-12 14:32:26 +0000931
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700932/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000933
934typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000935 PyObject_HEAD
936 PyObject *it;
937 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700938 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000939 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000940} cycleobject;
941
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000942static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000943
Tal Einatc4bccd32018-09-12 00:49:13 +0300944/*[clinic input]
945@classmethod
946itertools.cycle.__new__
947 iterable: object
948 /
949Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
950[clinic start generated code]*/
951
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000952static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300953itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
954/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000955{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000956 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000957 PyObject *saved;
958 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000959
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000960 /* Get iterator. */
961 it = PyObject_GetIter(iterable);
962 if (it == NULL)
963 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000964
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000965 saved = PyList_New(0);
966 if (saved == NULL) {
967 Py_DECREF(it);
968 return NULL;
969 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000970
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000971 /* create cycleobject structure */
972 lz = (cycleobject *)type->tp_alloc(type, 0);
973 if (lz == NULL) {
974 Py_DECREF(it);
975 Py_DECREF(saved);
976 return NULL;
977 }
978 lz->it = it;
979 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700980 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000981 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000982
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000983 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000984}
985
986static void
987cycle_dealloc(cycleobject *lz)
988{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000989 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000990 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700991 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000992 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000993}
994
995static int
996cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
997{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700998 if (lz->it)
999 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 Py_VISIT(lz->saved);
1001 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001002}
1003
1004static PyObject *
1005cycle_next(cycleobject *lz)
1006{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001007 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001008
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001009 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001010 item = PyIter_Next(lz->it);
1011 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001012 if (lz->firstpass)
1013 return item;
1014 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001015 Py_DECREF(item);
1016 return NULL;
1017 }
1018 return item;
1019 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001020 /* Note: StopIteration is already cleared by PyIter_Next() */
1021 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001022 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001023 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001024 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001025 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001026 return NULL;
1027 item = PyList_GET_ITEM(lz->saved, lz->index);
1028 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001029 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001030 lz->index = 0;
1031 Py_INCREF(item);
1032 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001033}
1034
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001035static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301036cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001037{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001038 /* Create a new cycle with the iterator tuple, then set the saved state */
1039 if (lz->it == NULL) {
1040 PyObject *it = PyObject_GetIter(lz->saved);
1041 if (it == NULL)
1042 return NULL;
1043 if (lz->index != 0) {
1044 _Py_IDENTIFIER(__setstate__);
1045 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1046 "n", lz->index);
1047 if (res == NULL) {
1048 Py_DECREF(it);
1049 return NULL;
1050 }
1051 Py_DECREF(res);
1052 }
1053 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
1054 }
1055 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
1056 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001057}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001058
1059static PyObject *
1060cycle_setstate(cycleobject *lz, PyObject *state)
1061{
1062 PyObject *saved=NULL;
1063 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001064 if (!PyTuple_Check(state)) {
1065 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001066 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001067 }
1068 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1069 return NULL;
1070 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001071 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001072 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001073 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001074 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001075 Py_RETURN_NONE;
1076}
1077
1078static PyMethodDef cycle_methods[] = {
1079 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1080 reduce_doc},
1081 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1082 setstate_doc},
1083 {NULL, NULL} /* sentinel */
1084};
1085
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001086static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001087 PyVarObject_HEAD_INIT(NULL, 0)
1088 "itertools.cycle", /* tp_name */
1089 sizeof(cycleobject), /* tp_basicsize */
1090 0, /* tp_itemsize */
1091 /* methods */
1092 (destructor)cycle_dealloc, /* tp_dealloc */
1093 0, /* tp_print */
1094 0, /* tp_getattr */
1095 0, /* tp_setattr */
1096 0, /* tp_reserved */
1097 0, /* tp_repr */
1098 0, /* tp_as_number */
1099 0, /* tp_as_sequence */
1100 0, /* tp_as_mapping */
1101 0, /* tp_hash */
1102 0, /* tp_call */
1103 0, /* tp_str */
1104 PyObject_GenericGetAttr, /* tp_getattro */
1105 0, /* tp_setattro */
1106 0, /* tp_as_buffer */
1107 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1108 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001109 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001110 (traverseproc)cycle_traverse, /* tp_traverse */
1111 0, /* tp_clear */
1112 0, /* tp_richcompare */
1113 0, /* tp_weaklistoffset */
1114 PyObject_SelfIter, /* tp_iter */
1115 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001116 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001117 0, /* tp_members */
1118 0, /* tp_getset */
1119 0, /* tp_base */
1120 0, /* tp_dict */
1121 0, /* tp_descr_get */
1122 0, /* tp_descr_set */
1123 0, /* tp_dictoffset */
1124 0, /* tp_init */
1125 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001126 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001127 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001128};
1129
1130
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131/* dropwhile object **********************************************************/
1132
1133typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 PyObject_HEAD
1135 PyObject *func;
1136 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001137 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001138} dropwhileobject;
1139
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001140static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001141
Tal Einatc4bccd32018-09-12 00:49:13 +03001142/*[clinic input]
1143@classmethod
1144itertools.dropwhile.__new__
1145 predicate as func: object
1146 iterable as seq: object
1147 /
1148Drop items from the iterable while predicate(item) is true.
1149
1150Afterwards, return every element until the iterable is exhausted.
1151[clinic start generated code]*/
1152
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001154itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1155/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001156{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001157 PyObject *it;
1158 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001160 /* Get iterator. */
1161 it = PyObject_GetIter(seq);
1162 if (it == NULL)
1163 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001164
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001165 /* create dropwhileobject structure */
1166 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1167 if (lz == NULL) {
1168 Py_DECREF(it);
1169 return NULL;
1170 }
1171 Py_INCREF(func);
1172 lz->func = func;
1173 lz->it = it;
1174 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001176 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177}
1178
1179static void
1180dropwhile_dealloc(dropwhileobject *lz)
1181{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001182 PyObject_GC_UnTrack(lz);
1183 Py_XDECREF(lz->func);
1184 Py_XDECREF(lz->it);
1185 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001186}
1187
1188static int
1189dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1190{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001191 Py_VISIT(lz->it);
1192 Py_VISIT(lz->func);
1193 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194}
1195
1196static PyObject *
1197dropwhile_next(dropwhileobject *lz)
1198{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001199 PyObject *item, *good;
1200 PyObject *it = lz->it;
1201 long ok;
1202 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001203
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001204 iternext = *Py_TYPE(it)->tp_iternext;
1205 for (;;) {
1206 item = iternext(it);
1207 if (item == NULL)
1208 return NULL;
1209 if (lz->start == 1)
1210 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001212 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001213 if (good == NULL) {
1214 Py_DECREF(item);
1215 return NULL;
1216 }
1217 ok = PyObject_IsTrue(good);
1218 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001219 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001220 lz->start = 1;
1221 return item;
1222 }
1223 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001224 if (ok < 0)
1225 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001226 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001227}
1228
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001229static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301230dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001231{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001232 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 +00001233}
1234
1235static PyObject *
1236dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1237{
1238 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001239 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001240 return NULL;
1241 lz->start = start;
1242 Py_RETURN_NONE;
1243}
1244
1245static PyMethodDef dropwhile_methods[] = {
1246 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1247 reduce_doc},
1248 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1249 setstate_doc},
1250 {NULL, NULL} /* sentinel */
1251};
1252
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001253static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001254 PyVarObject_HEAD_INIT(NULL, 0)
1255 "itertools.dropwhile", /* tp_name */
1256 sizeof(dropwhileobject), /* tp_basicsize */
1257 0, /* tp_itemsize */
1258 /* methods */
1259 (destructor)dropwhile_dealloc, /* tp_dealloc */
1260 0, /* tp_print */
1261 0, /* tp_getattr */
1262 0, /* tp_setattr */
1263 0, /* tp_reserved */
1264 0, /* tp_repr */
1265 0, /* tp_as_number */
1266 0, /* tp_as_sequence */
1267 0, /* tp_as_mapping */
1268 0, /* tp_hash */
1269 0, /* tp_call */
1270 0, /* tp_str */
1271 PyObject_GenericGetAttr, /* tp_getattro */
1272 0, /* tp_setattro */
1273 0, /* tp_as_buffer */
1274 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1275 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001276 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001277 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001278 0, /* tp_clear */
1279 0, /* tp_richcompare */
1280 0, /* tp_weaklistoffset */
1281 PyObject_SelfIter, /* tp_iter */
1282 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001283 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001284 0, /* tp_members */
1285 0, /* tp_getset */
1286 0, /* tp_base */
1287 0, /* tp_dict */
1288 0, /* tp_descr_get */
1289 0, /* tp_descr_set */
1290 0, /* tp_dictoffset */
1291 0, /* tp_init */
1292 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001293 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001294 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295};
1296
1297
1298/* takewhile object **********************************************************/
1299
1300typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001301 PyObject_HEAD
1302 PyObject *func;
1303 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001304 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001305} takewhileobject;
1306
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001307static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001308
Tal Einatc4bccd32018-09-12 00:49:13 +03001309/*[clinic input]
1310@classmethod
1311itertools.takewhile.__new__
1312 predicate as func: object
1313 iterable as seq: object
1314 /
1315Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1316[clinic start generated code]*/
1317
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001319itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1320/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001322 PyObject *it;
1323 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001325 /* Get iterator. */
1326 it = PyObject_GetIter(seq);
1327 if (it == NULL)
1328 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001329
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001330 /* create takewhileobject structure */
1331 lz = (takewhileobject *)type->tp_alloc(type, 0);
1332 if (lz == NULL) {
1333 Py_DECREF(it);
1334 return NULL;
1335 }
1336 Py_INCREF(func);
1337 lz->func = func;
1338 lz->it = it;
1339 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001341 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342}
1343
1344static void
1345takewhile_dealloc(takewhileobject *lz)
1346{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001347 PyObject_GC_UnTrack(lz);
1348 Py_XDECREF(lz->func);
1349 Py_XDECREF(lz->it);
1350 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351}
1352
1353static int
1354takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1355{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001356 Py_VISIT(lz->it);
1357 Py_VISIT(lz->func);
1358 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359}
1360
1361static PyObject *
1362takewhile_next(takewhileobject *lz)
1363{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001364 PyObject *item, *good;
1365 PyObject *it = lz->it;
1366 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001368 if (lz->stop == 1)
1369 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001371 item = (*Py_TYPE(it)->tp_iternext)(it);
1372 if (item == NULL)
1373 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001375 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001376 if (good == NULL) {
1377 Py_DECREF(item);
1378 return NULL;
1379 }
1380 ok = PyObject_IsTrue(good);
1381 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001382 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001383 return item;
1384 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001385 if (ok == 0)
1386 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001387 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388}
1389
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001390static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301391takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001392{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001393 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 +00001394}
1395
1396static PyObject *
1397takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1398{
1399 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001400
Antoine Pitrou721738f2012-08-15 23:20:39 +02001401 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001402 return NULL;
1403 lz->stop = stop;
1404 Py_RETURN_NONE;
1405}
1406
1407static PyMethodDef takewhile_reduce_methods[] = {
1408 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1409 reduce_doc},
1410 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1411 setstate_doc},
1412 {NULL, NULL} /* sentinel */
1413};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001414
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001415static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001416 PyVarObject_HEAD_INIT(NULL, 0)
1417 "itertools.takewhile", /* tp_name */
1418 sizeof(takewhileobject), /* tp_basicsize */
1419 0, /* tp_itemsize */
1420 /* methods */
1421 (destructor)takewhile_dealloc, /* tp_dealloc */
1422 0, /* tp_print */
1423 0, /* tp_getattr */
1424 0, /* tp_setattr */
1425 0, /* tp_reserved */
1426 0, /* tp_repr */
1427 0, /* tp_as_number */
1428 0, /* tp_as_sequence */
1429 0, /* tp_as_mapping */
1430 0, /* tp_hash */
1431 0, /* tp_call */
1432 0, /* tp_str */
1433 PyObject_GenericGetAttr, /* tp_getattro */
1434 0, /* tp_setattro */
1435 0, /* tp_as_buffer */
1436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1437 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001438 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001439 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001440 0, /* tp_clear */
1441 0, /* tp_richcompare */
1442 0, /* tp_weaklistoffset */
1443 PyObject_SelfIter, /* tp_iter */
1444 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001445 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001446 0, /* tp_members */
1447 0, /* tp_getset */
1448 0, /* tp_base */
1449 0, /* tp_dict */
1450 0, /* tp_descr_get */
1451 0, /* tp_descr_set */
1452 0, /* tp_dictoffset */
1453 0, /* tp_init */
1454 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001455 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457};
1458
1459
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001460/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
1462typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 PyObject_HEAD
1464 PyObject *it;
1465 Py_ssize_t next;
1466 Py_ssize_t stop;
1467 Py_ssize_t step;
1468 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469} isliceobject;
1470
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001471static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472
1473static PyObject *
1474islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1475{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001476 PyObject *seq;
1477 Py_ssize_t start=0, stop=-1, step=1;
1478 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1479 Py_ssize_t numargs;
1480 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001482 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001483 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001484
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001485 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1486 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001488 numargs = PyTuple_Size(args);
1489 if (numargs == 2) {
1490 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001491 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 if (stop == -1) {
1493 if (PyErr_Occurred())
1494 PyErr_Clear();
1495 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001496 "Stop argument for islice() must be None or "
1497 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 return NULL;
1499 }
1500 }
1501 } else {
1502 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001503 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 if (start == -1 && PyErr_Occurred())
1505 PyErr_Clear();
1506 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001507 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 if (stop == -1) {
1509 if (PyErr_Occurred())
1510 PyErr_Clear();
1511 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001512 "Stop argument for islice() must be None or "
1513 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001514 return NULL;
1515 }
1516 }
1517 }
1518 if (start<0 || stop<-1) {
1519 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001520 "Indices for islice() must be None or "
1521 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001522 return NULL;
1523 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001525 if (a3 != NULL) {
1526 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001527 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001528 if (step == -1 && PyErr_Occurred())
1529 PyErr_Clear();
1530 }
1531 if (step<1) {
1532 PyErr_SetString(PyExc_ValueError,
1533 "Step for islice() must be a positive integer or None.");
1534 return NULL;
1535 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001537 /* Get iterator. */
1538 it = PyObject_GetIter(seq);
1539 if (it == NULL)
1540 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001542 /* create isliceobject structure */
1543 lz = (isliceobject *)type->tp_alloc(type, 0);
1544 if (lz == NULL) {
1545 Py_DECREF(it);
1546 return NULL;
1547 }
1548 lz->it = it;
1549 lz->next = start;
1550 lz->stop = stop;
1551 lz->step = step;
1552 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001554 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001555}
1556
1557static void
1558islice_dealloc(isliceobject *lz)
1559{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001560 PyObject_GC_UnTrack(lz);
1561 Py_XDECREF(lz->it);
1562 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563}
1564
1565static int
1566islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1567{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001568 Py_VISIT(lz->it);
1569 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001570}
1571
1572static PyObject *
1573islice_next(isliceobject *lz)
1574{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001575 PyObject *item;
1576 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001577 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001578 Py_ssize_t oldnext;
1579 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001581 if (it == NULL)
1582 return NULL;
1583
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001584 iternext = *Py_TYPE(it)->tp_iternext;
1585 while (lz->cnt < lz->next) {
1586 item = iternext(it);
1587 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001588 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001589 Py_DECREF(item);
1590 lz->cnt++;
1591 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001592 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001593 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001594 item = iternext(it);
1595 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001596 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001597 lz->cnt++;
1598 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001599 /* The (size_t) cast below avoids the danger of undefined
1600 behaviour from signed integer overflow. */
1601 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001602 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1603 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001604 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001605
1606empty:
1607 Py_CLEAR(lz->it);
1608 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609}
1610
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001611static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301612islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001613{
1614 /* When unpickled, generate a new object with the same bounds,
1615 * then 'setstate' with the next and count
1616 */
1617 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001618
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001619 if (lz->it == NULL) {
1620 PyObject *empty_list;
1621 PyObject *empty_it;
1622 empty_list = PyList_New(0);
1623 if (empty_list == NULL)
1624 return NULL;
1625 empty_it = PyObject_GetIter(empty_list);
1626 Py_DECREF(empty_list);
1627 if (empty_it == NULL)
1628 return NULL;
1629 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1630 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001631 if (lz->stop == -1) {
1632 stop = Py_None;
1633 Py_INCREF(stop);
1634 } else {
1635 stop = PyLong_FromSsize_t(lz->stop);
1636 if (stop == NULL)
1637 return NULL;
1638 }
1639 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1640 lz->it, lz->next, stop, lz->step,
1641 lz->cnt);
1642}
1643
1644static PyObject *
1645islice_setstate(isliceobject *lz, PyObject *state)
1646{
1647 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001648
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001649 if (cnt == -1 && PyErr_Occurred())
1650 return NULL;
1651 lz->cnt = cnt;
1652 Py_RETURN_NONE;
1653}
1654
1655static PyMethodDef islice_methods[] = {
1656 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1657 reduce_doc},
1658 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1659 setstate_doc},
1660 {NULL, NULL} /* sentinel */
1661};
1662
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001663PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001664"islice(iterable, stop) --> islice object\n\
1665islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666\n\
1667Return an iterator whose next() method returns selected values from an\n\
1668iterable. If start is specified, will skip all preceding elements;\n\
1669otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001670specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671skipped between successive calls. Works like a slice() on a list\n\
1672but returns an iterator.");
1673
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001674static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001675 PyVarObject_HEAD_INIT(NULL, 0)
1676 "itertools.islice", /* tp_name */
1677 sizeof(isliceobject), /* tp_basicsize */
1678 0, /* tp_itemsize */
1679 /* methods */
1680 (destructor)islice_dealloc, /* tp_dealloc */
1681 0, /* tp_print */
1682 0, /* tp_getattr */
1683 0, /* tp_setattr */
1684 0, /* tp_reserved */
1685 0, /* tp_repr */
1686 0, /* tp_as_number */
1687 0, /* tp_as_sequence */
1688 0, /* tp_as_mapping */
1689 0, /* tp_hash */
1690 0, /* tp_call */
1691 0, /* tp_str */
1692 PyObject_GenericGetAttr, /* tp_getattro */
1693 0, /* tp_setattro */
1694 0, /* tp_as_buffer */
1695 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1696 Py_TPFLAGS_BASETYPE, /* tp_flags */
1697 islice_doc, /* tp_doc */
1698 (traverseproc)islice_traverse, /* tp_traverse */
1699 0, /* tp_clear */
1700 0, /* tp_richcompare */
1701 0, /* tp_weaklistoffset */
1702 PyObject_SelfIter, /* tp_iter */
1703 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001704 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001705 0, /* tp_members */
1706 0, /* tp_getset */
1707 0, /* tp_base */
1708 0, /* tp_dict */
1709 0, /* tp_descr_get */
1710 0, /* tp_descr_set */
1711 0, /* tp_dictoffset */
1712 0, /* tp_init */
1713 0, /* tp_alloc */
1714 islice_new, /* tp_new */
1715 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716};
1717
1718
1719/* starmap object ************************************************************/
1720
1721typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001722 PyObject_HEAD
1723 PyObject *func;
1724 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001725} starmapobject;
1726
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001727static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001728
Tal Einatc4bccd32018-09-12 00:49:13 +03001729/*[clinic input]
1730@classmethod
1731itertools.starmap.__new__
1732 function as func: object
1733 iterable as seq: object
1734 /
1735Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1736[clinic start generated code]*/
1737
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001738static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001739itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1740/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001741{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001742 PyObject *it;
1743 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001745 /* Get iterator. */
1746 it = PyObject_GetIter(seq);
1747 if (it == NULL)
1748 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001750 /* create starmapobject structure */
1751 lz = (starmapobject *)type->tp_alloc(type, 0);
1752 if (lz == NULL) {
1753 Py_DECREF(it);
1754 return NULL;
1755 }
1756 Py_INCREF(func);
1757 lz->func = func;
1758 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001760 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001761}
1762
1763static void
1764starmap_dealloc(starmapobject *lz)
1765{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001766 PyObject_GC_UnTrack(lz);
1767 Py_XDECREF(lz->func);
1768 Py_XDECREF(lz->it);
1769 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001770}
1771
1772static int
1773starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1774{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001775 Py_VISIT(lz->it);
1776 Py_VISIT(lz->func);
1777 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001778}
1779
1780static PyObject *
1781starmap_next(starmapobject *lz)
1782{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001783 PyObject *args;
1784 PyObject *result;
1785 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001786
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001787 args = (*Py_TYPE(it)->tp_iternext)(it);
1788 if (args == NULL)
1789 return NULL;
1790 if (!PyTuple_CheckExact(args)) {
1791 PyObject *newargs = PySequence_Tuple(args);
1792 Py_DECREF(args);
1793 if (newargs == NULL)
1794 return NULL;
1795 args = newargs;
1796 }
1797 result = PyObject_Call(lz->func, args, NULL);
1798 Py_DECREF(args);
1799 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001800}
1801
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001802static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301803starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001804{
1805 /* Just pickle the iterator */
1806 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1807}
1808
1809static PyMethodDef starmap_methods[] = {
1810 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1811 reduce_doc},
1812 {NULL, NULL} /* sentinel */
1813};
1814
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001815static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001816 PyVarObject_HEAD_INIT(NULL, 0)
1817 "itertools.starmap", /* tp_name */
1818 sizeof(starmapobject), /* tp_basicsize */
1819 0, /* tp_itemsize */
1820 /* methods */
1821 (destructor)starmap_dealloc, /* tp_dealloc */
1822 0, /* tp_print */
1823 0, /* tp_getattr */
1824 0, /* tp_setattr */
1825 0, /* tp_reserved */
1826 0, /* tp_repr */
1827 0, /* tp_as_number */
1828 0, /* tp_as_sequence */
1829 0, /* tp_as_mapping */
1830 0, /* tp_hash */
1831 0, /* tp_call */
1832 0, /* tp_str */
1833 PyObject_GenericGetAttr, /* tp_getattro */
1834 0, /* tp_setattro */
1835 0, /* tp_as_buffer */
1836 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1837 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001838 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001839 (traverseproc)starmap_traverse, /* tp_traverse */
1840 0, /* tp_clear */
1841 0, /* tp_richcompare */
1842 0, /* tp_weaklistoffset */
1843 PyObject_SelfIter, /* tp_iter */
1844 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001845 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001846 0, /* tp_members */
1847 0, /* tp_getset */
1848 0, /* tp_base */
1849 0, /* tp_dict */
1850 0, /* tp_descr_get */
1851 0, /* tp_descr_set */
1852 0, /* tp_dictoffset */
1853 0, /* tp_init */
1854 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001855 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001856 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001857};
1858
1859
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001860/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001861
1862typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001863 PyObject_HEAD
1864 PyObject *source; /* Iterator over input iterables */
1865 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001866} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001868static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001869
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001870static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001871chain_new_internal(PyTypeObject *type, PyObject *source)
1872{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001873 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001875 lz = (chainobject *)type->tp_alloc(type, 0);
1876 if (lz == NULL) {
1877 Py_DECREF(source);
1878 return NULL;
1879 }
1880
1881 lz->source = source;
1882 lz->active = NULL;
1883 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001884}
1885
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001886static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001887chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001888{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001889 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001890
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001891 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001892 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001893
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001894 source = PyObject_GetIter(args);
1895 if (source == NULL)
1896 return NULL;
1897
1898 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001899}
1900
Tal Einatc4bccd32018-09-12 00:49:13 +03001901/*[clinic input]
1902@classmethod
1903itertools.chain.from_iterable
1904 iterable as arg: object
1905 /
1906Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1907[clinic start generated code]*/
1908
Christian Heimesf16baeb2008-02-29 14:57:44 +00001909static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001910itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1911/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001913 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001914
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001915 source = PyObject_GetIter(arg);
1916 if (source == NULL)
1917 return NULL;
1918
1919 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001920}
1921
1922static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001923chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001924{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001925 PyObject_GC_UnTrack(lz);
1926 Py_XDECREF(lz->active);
1927 Py_XDECREF(lz->source);
1928 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001929}
1930
Raymond Hettinger2012f172003-02-07 05:32:58 +00001931static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001932chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001933{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001934 Py_VISIT(lz->source);
1935 Py_VISIT(lz->active);
1936 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001937}
1938
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001939static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001940chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001941{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001942 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001943
T. Wouters5466d4a2017-03-30 09:58:35 -07001944 /* lz->source is the iterator of iterables. If it's NULL, we've already
1945 * consumed them all. lz->active is the current iterator. If it's NULL,
1946 * we should grab a new one from lz->source. */
1947 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001948 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001949 PyObject *iterable = PyIter_Next(lz->source);
1950 if (iterable == NULL) {
1951 Py_CLEAR(lz->source);
1952 return NULL; /* no more input sources */
1953 }
1954 lz->active = PyObject_GetIter(iterable);
1955 Py_DECREF(iterable);
1956 if (lz->active == NULL) {
1957 Py_CLEAR(lz->source);
1958 return NULL; /* input not iterable */
1959 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001960 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001961 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1962 if (item != NULL)
1963 return item;
1964 if (PyErr_Occurred()) {
1965 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1966 PyErr_Clear();
1967 else
1968 return NULL; /* input raised an exception */
1969 }
1970 /* lz->active is consumed, try with the next iterable. */
1971 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001972 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001973 /* Everything had been consumed already. */
1974 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001975}
1976
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001977static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301978chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001979{
1980 if (lz->source) {
1981 /* we can't pickle function objects (itertools.from_iterable) so
1982 * we must use setstate to replace the iterable. One day we
1983 * will fix pickling of functions
1984 */
1985 if (lz->active) {
1986 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1987 } else {
1988 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1989 }
1990 } else {
1991 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1992 }
1993 return NULL;
1994}
1995
1996static PyObject *
1997chain_setstate(chainobject *lz, PyObject *state)
1998{
1999 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002000
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002001 if (!PyTuple_Check(state)) {
2002 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002003 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002004 }
2005 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2006 return NULL;
2007 }
2008 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2009 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2010 return NULL;
2011 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002012
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002013 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002014 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002015 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002016 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002017 Py_RETURN_NONE;
2018}
2019
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002020PyDoc_STRVAR(chain_doc,
2021"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002022\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002023Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002024first iterable until it is exhausted, then elements from the next\n\
2025iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002026
Christian Heimesf16baeb2008-02-29 14:57:44 +00002027static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002028 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002029 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2030 reduce_doc},
2031 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2032 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002033 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002034};
2035
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002036static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002037 PyVarObject_HEAD_INIT(NULL, 0)
2038 "itertools.chain", /* tp_name */
2039 sizeof(chainobject), /* tp_basicsize */
2040 0, /* tp_itemsize */
2041 /* methods */
2042 (destructor)chain_dealloc, /* tp_dealloc */
2043 0, /* tp_print */
2044 0, /* tp_getattr */
2045 0, /* tp_setattr */
2046 0, /* tp_reserved */
2047 0, /* tp_repr */
2048 0, /* tp_as_number */
2049 0, /* tp_as_sequence */
2050 0, /* tp_as_mapping */
2051 0, /* tp_hash */
2052 0, /* tp_call */
2053 0, /* tp_str */
2054 PyObject_GenericGetAttr, /* tp_getattro */
2055 0, /* tp_setattro */
2056 0, /* tp_as_buffer */
2057 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2058 Py_TPFLAGS_BASETYPE, /* tp_flags */
2059 chain_doc, /* tp_doc */
2060 (traverseproc)chain_traverse, /* tp_traverse */
2061 0, /* tp_clear */
2062 0, /* tp_richcompare */
2063 0, /* tp_weaklistoffset */
2064 PyObject_SelfIter, /* tp_iter */
2065 (iternextfunc)chain_next, /* tp_iternext */
2066 chain_methods, /* tp_methods */
2067 0, /* tp_members */
2068 0, /* tp_getset */
2069 0, /* tp_base */
2070 0, /* tp_dict */
2071 0, /* tp_descr_get */
2072 0, /* tp_descr_set */
2073 0, /* tp_dictoffset */
2074 0, /* tp_init */
2075 0, /* tp_alloc */
2076 chain_new, /* tp_new */
2077 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002078};
2079
2080
Christian Heimesc3f30c42008-02-22 16:37:40 +00002081/* product object ************************************************************/
2082
2083typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002084 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002085 PyObject *pools; /* tuple of pool tuples */
2086 Py_ssize_t *indices; /* one index per pool */
2087 PyObject *result; /* most recently returned result tuple */
2088 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002089} productobject;
2090
2091static PyTypeObject product_type;
2092
2093static PyObject *
2094product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2095{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002096 productobject *lz;
2097 Py_ssize_t nargs, npools, repeat=1;
2098 PyObject *pools = NULL;
2099 Py_ssize_t *indices = NULL;
2100 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002102 if (kwds != NULL) {
2103 char *kwlist[] = {"repeat", 0};
2104 PyObject *tmpargs = PyTuple_New(0);
2105 if (tmpargs == NULL)
2106 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002107 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2108 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002109 Py_DECREF(tmpargs);
2110 return NULL;
2111 }
2112 Py_DECREF(tmpargs);
2113 if (repeat < 0) {
2114 PyErr_SetString(PyExc_ValueError,
2115 "repeat argument cannot be negative");
2116 return NULL;
2117 }
2118 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002119
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002120 assert(PyTuple_CheckExact(args));
2121 if (repeat == 0) {
2122 nargs = 0;
2123 } else {
2124 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002125 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002126 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2127 return NULL;
2128 }
2129 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002130 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002131
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002132 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002133 if (indices == NULL) {
2134 PyErr_NoMemory();
2135 goto error;
2136 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002137
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002138 pools = PyTuple_New(npools);
2139 if (pools == NULL)
2140 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002141
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002142 for (i=0; i < nargs ; ++i) {
2143 PyObject *item = PyTuple_GET_ITEM(args, i);
2144 PyObject *pool = PySequence_Tuple(item);
2145 if (pool == NULL)
2146 goto error;
2147 PyTuple_SET_ITEM(pools, i, pool);
2148 indices[i] = 0;
2149 }
2150 for ( ; i < npools; ++i) {
2151 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2152 Py_INCREF(pool);
2153 PyTuple_SET_ITEM(pools, i, pool);
2154 indices[i] = 0;
2155 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002156
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002157 /* create productobject structure */
2158 lz = (productobject *)type->tp_alloc(type, 0);
2159 if (lz == NULL)
2160 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002161
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002162 lz->pools = pools;
2163 lz->indices = indices;
2164 lz->result = NULL;
2165 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002167 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002168
2169error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002170 if (indices != NULL)
2171 PyMem_Free(indices);
2172 Py_XDECREF(pools);
2173 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002174}
2175
2176static void
2177product_dealloc(productobject *lz)
2178{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002179 PyObject_GC_UnTrack(lz);
2180 Py_XDECREF(lz->pools);
2181 Py_XDECREF(lz->result);
2182 if (lz->indices != NULL)
2183 PyMem_Free(lz->indices);
2184 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002185}
2186
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002187static PyObject *
2188product_sizeof(productobject *lz, void *unused)
2189{
2190 Py_ssize_t res;
2191
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002192 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002193 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2194 return PyLong_FromSsize_t(res);
2195}
2196
2197PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2198
Christian Heimesc3f30c42008-02-22 16:37:40 +00002199static int
2200product_traverse(productobject *lz, visitproc visit, void *arg)
2201{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002202 Py_VISIT(lz->pools);
2203 Py_VISIT(lz->result);
2204 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002205}
2206
2207static PyObject *
2208product_next(productobject *lz)
2209{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002210 PyObject *pool;
2211 PyObject *elem;
2212 PyObject *oldelem;
2213 PyObject *pools = lz->pools;
2214 PyObject *result = lz->result;
2215 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2216 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002217
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002218 if (lz->stopped)
2219 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002220
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002221 if (result == NULL) {
2222 /* On the first pass, return an initial tuple filled with the
2223 first element from each pool. */
2224 result = PyTuple_New(npools);
2225 if (result == NULL)
2226 goto empty;
2227 lz->result = result;
2228 for (i=0; i < npools; i++) {
2229 pool = PyTuple_GET_ITEM(pools, i);
2230 if (PyTuple_GET_SIZE(pool) == 0)
2231 goto empty;
2232 elem = PyTuple_GET_ITEM(pool, 0);
2233 Py_INCREF(elem);
2234 PyTuple_SET_ITEM(result, i, elem);
2235 }
2236 } else {
2237 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002238
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002239 /* Copy the previous result tuple or re-use it if available */
2240 if (Py_REFCNT(result) > 1) {
2241 PyObject *old_result = result;
2242 result = PyTuple_New(npools);
2243 if (result == NULL)
2244 goto empty;
2245 lz->result = result;
2246 for (i=0; i < npools; i++) {
2247 elem = PyTuple_GET_ITEM(old_result, i);
2248 Py_INCREF(elem);
2249 PyTuple_SET_ITEM(result, i, elem);
2250 }
2251 Py_DECREF(old_result);
2252 }
2253 /* Now, we've got the only copy so we can update it in-place */
2254 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002255
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002256 /* Update the pool indices right-to-left. Only advance to the
2257 next pool when the previous one rolls-over */
2258 for (i=npools-1 ; i >= 0 ; i--) {
2259 pool = PyTuple_GET_ITEM(pools, i);
2260 indices[i]++;
2261 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2262 /* Roll-over and advance to next pool */
2263 indices[i] = 0;
2264 elem = PyTuple_GET_ITEM(pool, 0);
2265 Py_INCREF(elem);
2266 oldelem = PyTuple_GET_ITEM(result, i);
2267 PyTuple_SET_ITEM(result, i, elem);
2268 Py_DECREF(oldelem);
2269 } else {
2270 /* No rollover. Just increment and stop here. */
2271 elem = PyTuple_GET_ITEM(pool, indices[i]);
2272 Py_INCREF(elem);
2273 oldelem = PyTuple_GET_ITEM(result, i);
2274 PyTuple_SET_ITEM(result, i, elem);
2275 Py_DECREF(oldelem);
2276 break;
2277 }
2278 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002279
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002280 /* If i is negative, then the indices have all rolled-over
2281 and we're done. */
2282 if (i < 0)
2283 goto empty;
2284 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002285
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002286 Py_INCREF(result);
2287 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002288
2289empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002290 lz->stopped = 1;
2291 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002292}
2293
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002294static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302295product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002296{
2297 if (lz->stopped) {
2298 return Py_BuildValue("O(())", Py_TYPE(lz));
2299 } else if (lz->result == NULL) {
2300 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2301 } else {
2302 PyObject *indices;
2303 Py_ssize_t n, i;
2304
2305 /* we must pickle the indices use them for setstate, and
2306 * additionally indicate that the iterator has started
2307 */
2308 n = PyTuple_GET_SIZE(lz->pools);
2309 indices = PyTuple_New(n);
2310 if (indices == NULL)
2311 return NULL;
2312 for (i=0; i<n; i++){
2313 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2314 if (!index) {
2315 Py_DECREF(indices);
2316 return NULL;
2317 }
2318 PyTuple_SET_ITEM(indices, i, index);
2319 }
2320 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2321 }
2322}
2323
2324static PyObject *
2325product_setstate(productobject *lz, PyObject *state)
2326{
2327 PyObject *result;
2328 Py_ssize_t n, i;
2329
2330 n = PyTuple_GET_SIZE(lz->pools);
2331 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2332 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2333 return NULL;
2334 }
2335 for (i=0; i<n; i++)
2336 {
2337 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2338 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002339 PyObject* pool;
2340 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002341 if (index < 0 && PyErr_Occurred())
2342 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002343 pool = PyTuple_GET_ITEM(lz->pools, i);
2344 poolsize = PyTuple_GET_SIZE(pool);
2345 if (poolsize == 0) {
2346 lz->stopped = 1;
2347 Py_RETURN_NONE;
2348 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002349 /* clamp the index */
2350 if (index < 0)
2351 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002352 else if (index > poolsize-1)
2353 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002354 lz->indices[i] = index;
2355 }
2356
2357 result = PyTuple_New(n);
2358 if (!result)
2359 return NULL;
2360 for (i=0; i<n; i++) {
2361 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2362 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2363 Py_INCREF(element);
2364 PyTuple_SET_ITEM(result, i, element);
2365 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002366 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002367 Py_RETURN_NONE;
2368}
2369
2370static PyMethodDef product_methods[] = {
2371 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2372 reduce_doc},
2373 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2374 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002375 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2376 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002377 {NULL, NULL} /* sentinel */
2378};
2379
Christian Heimesc3f30c42008-02-22 16:37:40 +00002380PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002381"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002382\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002383Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002384For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2385The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2386cycle in a manner similar to an odometer (with the rightmost element changing\n\
2387on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002388To compute the product of an iterable with itself, specify the number\n\
2389of repetitions with the optional repeat keyword argument. For example,\n\
2390product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002391product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2392product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2393
2394static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002395 PyVarObject_HEAD_INIT(NULL, 0)
2396 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002397 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002398 0, /* tp_itemsize */
2399 /* methods */
2400 (destructor)product_dealloc, /* tp_dealloc */
2401 0, /* tp_print */
2402 0, /* tp_getattr */
2403 0, /* tp_setattr */
2404 0, /* tp_reserved */
2405 0, /* tp_repr */
2406 0, /* tp_as_number */
2407 0, /* tp_as_sequence */
2408 0, /* tp_as_mapping */
2409 0, /* tp_hash */
2410 0, /* tp_call */
2411 0, /* tp_str */
2412 PyObject_GenericGetAttr, /* tp_getattro */
2413 0, /* tp_setattro */
2414 0, /* tp_as_buffer */
2415 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2416 Py_TPFLAGS_BASETYPE, /* tp_flags */
2417 product_doc, /* tp_doc */
2418 (traverseproc)product_traverse, /* tp_traverse */
2419 0, /* tp_clear */
2420 0, /* tp_richcompare */
2421 0, /* tp_weaklistoffset */
2422 PyObject_SelfIter, /* tp_iter */
2423 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002424 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002425 0, /* tp_members */
2426 0, /* tp_getset */
2427 0, /* tp_base */
2428 0, /* tp_dict */
2429 0, /* tp_descr_get */
2430 0, /* tp_descr_set */
2431 0, /* tp_dictoffset */
2432 0, /* tp_init */
2433 0, /* tp_alloc */
2434 product_new, /* tp_new */
2435 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002436};
2437
2438
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002439/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002440
2441typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002442 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002443 PyObject *pool; /* input converted to a tuple */
2444 Py_ssize_t *indices; /* one index per result element */
2445 PyObject *result; /* most recently returned result tuple */
2446 Py_ssize_t r; /* size of result tuple */
2447 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002448} combinationsobject;
2449
2450static PyTypeObject combinations_type;
2451
Tal Einatc4bccd32018-09-12 00:49:13 +03002452
2453/*[clinic input]
2454@classmethod
2455itertools.combinations.__new__
2456 iterable: object
2457 r: Py_ssize_t
2458Return successive r-length combinations of elements in the iterable.
2459
2460combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2461[clinic start generated code]*/
2462
Christian Heimes380f7f22008-02-28 11:19:05 +00002463static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002464itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2465 Py_ssize_t r)
2466/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002467{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002468 combinationsobject *co;
2469 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002470 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002471 Py_ssize_t *indices = NULL;
2472 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002473
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002474 pool = PySequence_Tuple(iterable);
2475 if (pool == NULL)
2476 goto error;
2477 n = PyTuple_GET_SIZE(pool);
2478 if (r < 0) {
2479 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2480 goto error;
2481 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002482
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002483 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002484 if (indices == NULL) {
2485 PyErr_NoMemory();
2486 goto error;
2487 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002489 for (i=0 ; i<r ; i++)
2490 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002491
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002492 /* create combinationsobject structure */
2493 co = (combinationsobject *)type->tp_alloc(type, 0);
2494 if (co == NULL)
2495 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002496
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002497 co->pool = pool;
2498 co->indices = indices;
2499 co->result = NULL;
2500 co->r = r;
2501 co->stopped = r > n ? 1 : 0;
2502
2503 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002504
2505error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002506 if (indices != NULL)
2507 PyMem_Free(indices);
2508 Py_XDECREF(pool);
2509 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002510}
2511
2512static void
2513combinations_dealloc(combinationsobject *co)
2514{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002515 PyObject_GC_UnTrack(co);
2516 Py_XDECREF(co->pool);
2517 Py_XDECREF(co->result);
2518 if (co->indices != NULL)
2519 PyMem_Free(co->indices);
2520 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002521}
2522
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002523static PyObject *
2524combinations_sizeof(combinationsobject *co, void *unused)
2525{
2526 Py_ssize_t res;
2527
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002528 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002529 res += co->r * sizeof(Py_ssize_t);
2530 return PyLong_FromSsize_t(res);
2531}
2532
Christian Heimes380f7f22008-02-28 11:19:05 +00002533static int
2534combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2535{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002536 Py_VISIT(co->pool);
2537 Py_VISIT(co->result);
2538 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002539}
2540
2541static PyObject *
2542combinations_next(combinationsobject *co)
2543{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002544 PyObject *elem;
2545 PyObject *oldelem;
2546 PyObject *pool = co->pool;
2547 Py_ssize_t *indices = co->indices;
2548 PyObject *result = co->result;
2549 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2550 Py_ssize_t r = co->r;
2551 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002552
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002553 if (co->stopped)
2554 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002555
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002556 if (result == NULL) {
2557 /* On the first pass, initialize result tuple using the indices */
2558 result = PyTuple_New(r);
2559 if (result == NULL)
2560 goto empty;
2561 co->result = result;
2562 for (i=0; i<r ; i++) {
2563 index = indices[i];
2564 elem = PyTuple_GET_ITEM(pool, index);
2565 Py_INCREF(elem);
2566 PyTuple_SET_ITEM(result, i, elem);
2567 }
2568 } else {
2569 /* Copy the previous result tuple or re-use it if available */
2570 if (Py_REFCNT(result) > 1) {
2571 PyObject *old_result = result;
2572 result = PyTuple_New(r);
2573 if (result == NULL)
2574 goto empty;
2575 co->result = result;
2576 for (i=0; i<r ; i++) {
2577 elem = PyTuple_GET_ITEM(old_result, i);
2578 Py_INCREF(elem);
2579 PyTuple_SET_ITEM(result, i, elem);
2580 }
2581 Py_DECREF(old_result);
2582 }
2583 /* Now, we've got the only copy so we can update it in-place
2584 * CPython's empty tuple is a singleton and cached in
2585 * PyTuple's freelist.
2586 */
2587 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002588
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002589 /* Scan indices right-to-left until finding one that is not
2590 at its maximum (i + n - r). */
2591 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2592 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002593
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002594 /* If i is negative, then the indices are all at
2595 their maximum value and we're done. */
2596 if (i < 0)
2597 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002598
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002599 /* Increment the current index which we know is not at its
2600 maximum. Then move back to the right setting each index
2601 to its lowest possible value (one higher than the index
2602 to its left -- this maintains the sort order invariant). */
2603 indices[i]++;
2604 for (j=i+1 ; j<r ; j++)
2605 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002606
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002607 /* Update the result tuple for the new indices
2608 starting with i, the leftmost index that changed */
2609 for ( ; i<r ; i++) {
2610 index = indices[i];
2611 elem = PyTuple_GET_ITEM(pool, index);
2612 Py_INCREF(elem);
2613 oldelem = PyTuple_GET_ITEM(result, i);
2614 PyTuple_SET_ITEM(result, i, elem);
2615 Py_DECREF(oldelem);
2616 }
2617 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002618
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002619 Py_INCREF(result);
2620 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002621
2622empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002623 co->stopped = 1;
2624 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002625}
2626
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002627static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302628combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002629{
2630 if (lz->result == NULL) {
2631 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2632 } else if (lz->stopped) {
2633 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2634 } else {
2635 PyObject *indices;
2636 Py_ssize_t i;
2637
2638 /* we must pickle the indices and use them for setstate */
2639 indices = PyTuple_New(lz->r);
2640 if (!indices)
2641 return NULL;
2642 for (i=0; i<lz->r; i++)
2643 {
2644 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2645 if (!index) {
2646 Py_DECREF(indices);
2647 return NULL;
2648 }
2649 PyTuple_SET_ITEM(indices, i, index);
2650 }
2651
2652 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2653 }
2654}
2655
2656static PyObject *
2657combinations_setstate(combinationsobject *lz, PyObject *state)
2658{
2659 PyObject *result;
2660 Py_ssize_t i;
2661 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2662
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002663 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002664 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2665 return NULL;
2666 }
2667
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002668 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002669 Py_ssize_t max;
2670 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2671 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002672
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002673 if (index == -1 && PyErr_Occurred())
2674 return NULL; /* not an integer */
2675 max = i + n - lz->r;
2676 /* clamp the index (beware of negative max) */
2677 if (index > max)
2678 index = max;
2679 if (index < 0)
2680 index = 0;
2681 lz->indices[i] = index;
2682 }
2683
2684 result = PyTuple_New(lz->r);
2685 if (result == NULL)
2686 return NULL;
2687 for (i=0; i<lz->r; i++) {
2688 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2689 Py_INCREF(element);
2690 PyTuple_SET_ITEM(result, i, element);
2691 }
2692
Serhiy Storchaka48842712016-04-06 09:45:48 +03002693 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002694 Py_RETURN_NONE;
2695}
2696
2697static PyMethodDef combinations_methods[] = {
2698 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2699 reduce_doc},
2700 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2701 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002702 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2703 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002704 {NULL, NULL} /* sentinel */
2705};
2706
Christian Heimes380f7f22008-02-28 11:19:05 +00002707static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002709 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002710 sizeof(combinationsobject), /* tp_basicsize */
2711 0, /* tp_itemsize */
2712 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002713 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002714 0, /* tp_print */
2715 0, /* tp_getattr */
2716 0, /* tp_setattr */
2717 0, /* tp_reserved */
2718 0, /* tp_repr */
2719 0, /* tp_as_number */
2720 0, /* tp_as_sequence */
2721 0, /* tp_as_mapping */
2722 0, /* tp_hash */
2723 0, /* tp_call */
2724 0, /* tp_str */
2725 PyObject_GenericGetAttr, /* tp_getattro */
2726 0, /* tp_setattro */
2727 0, /* tp_as_buffer */
2728 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2729 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002730 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002731 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002732 0, /* tp_clear */
2733 0, /* tp_richcompare */
2734 0, /* tp_weaklistoffset */
2735 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002736 (iternextfunc)combinations_next, /* tp_iternext */
2737 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002738 0, /* tp_members */
2739 0, /* tp_getset */
2740 0, /* tp_base */
2741 0, /* tp_dict */
2742 0, /* tp_descr_get */
2743 0, /* tp_descr_set */
2744 0, /* tp_dictoffset */
2745 0, /* tp_init */
2746 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002747 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002749};
2750
2751
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002752/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002753
2754/* Equivalent to:
2755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 def combinations_with_replacement(iterable, r):
2757 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2758 # number items returned: (n+r-1)! / r! / (n-1)!
2759 pool = tuple(iterable)
2760 n = len(pool)
2761 indices = [0] * r
2762 yield tuple(pool[i] for i in indices)
2763 while 1:
2764 for i in reversed(range(r)):
2765 if indices[i] != n - 1:
2766 break
2767 else:
2768 return
2769 indices[i:] = [indices[i] + 1] * (r - i)
2770 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002771
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 def combinations_with_replacement2(iterable, r):
2773 'Alternate version that filters from product()'
2774 pool = tuple(iterable)
2775 n = len(pool)
2776 for indices in product(range(n), repeat=r):
2777 if sorted(indices) == list(indices):
2778 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002779*/
2780typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002781 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002782 PyObject *pool; /* input converted to a tuple */
2783 Py_ssize_t *indices; /* one index per result element */
2784 PyObject *result; /* most recently returned result tuple */
2785 Py_ssize_t r; /* size of result tuple */
2786 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002787} cwrobject;
2788
2789static PyTypeObject cwr_type;
2790
Tal Einatc4bccd32018-09-12 00:49:13 +03002791/*[clinic input]
2792@classmethod
2793itertools.combinations_with_replacement.__new__
2794 iterable: object
2795 r: Py_ssize_t
2796Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2797
2798combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2799[clinic start generated code]*/
2800
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002801static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002802itertools_combinations_with_replacement_impl(PyTypeObject *type,
2803 PyObject *iterable,
2804 Py_ssize_t r)
2805/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002806{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002807 cwrobject *co;
2808 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002810 Py_ssize_t *indices = NULL;
2811 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002813 pool = PySequence_Tuple(iterable);
2814 if (pool == NULL)
2815 goto error;
2816 n = PyTuple_GET_SIZE(pool);
2817 if (r < 0) {
2818 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2819 goto error;
2820 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002821
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002822 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002823 if (indices == NULL) {
2824 PyErr_NoMemory();
2825 goto error;
2826 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002827
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002828 for (i=0 ; i<r ; i++)
2829 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002831 /* create cwrobject structure */
2832 co = (cwrobject *)type->tp_alloc(type, 0);
2833 if (co == NULL)
2834 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002835
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002836 co->pool = pool;
2837 co->indices = indices;
2838 co->result = NULL;
2839 co->r = r;
2840 co->stopped = !n && r;
2841
2842 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002843
2844error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002845 if (indices != NULL)
2846 PyMem_Free(indices);
2847 Py_XDECREF(pool);
2848 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002849}
2850
2851static void
2852cwr_dealloc(cwrobject *co)
2853{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002854 PyObject_GC_UnTrack(co);
2855 Py_XDECREF(co->pool);
2856 Py_XDECREF(co->result);
2857 if (co->indices != NULL)
2858 PyMem_Free(co->indices);
2859 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002860}
2861
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002862static PyObject *
2863cwr_sizeof(cwrobject *co, void *unused)
2864{
2865 Py_ssize_t res;
2866
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002867 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002868 res += co->r * sizeof(Py_ssize_t);
2869 return PyLong_FromSsize_t(res);
2870}
2871
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002872static int
2873cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2874{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002875 Py_VISIT(co->pool);
2876 Py_VISIT(co->result);
2877 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002878}
2879
2880static PyObject *
2881cwr_next(cwrobject *co)
2882{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002883 PyObject *elem;
2884 PyObject *oldelem;
2885 PyObject *pool = co->pool;
2886 Py_ssize_t *indices = co->indices;
2887 PyObject *result = co->result;
2888 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2889 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002890 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002891
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002892 if (co->stopped)
2893 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002895 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002896 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002897 result = PyTuple_New(r);
2898 if (result == NULL)
2899 goto empty;
2900 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002901 if (n > 0) {
2902 elem = PyTuple_GET_ITEM(pool, 0);
2903 for (i=0; i<r ; i++) {
2904 assert(indices[i] == 0);
2905 Py_INCREF(elem);
2906 PyTuple_SET_ITEM(result, i, elem);
2907 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002908 }
2909 } else {
2910 /* Copy the previous result tuple or re-use it if available */
2911 if (Py_REFCNT(result) > 1) {
2912 PyObject *old_result = result;
2913 result = PyTuple_New(r);
2914 if (result == NULL)
2915 goto empty;
2916 co->result = result;
2917 for (i=0; i<r ; i++) {
2918 elem = PyTuple_GET_ITEM(old_result, i);
2919 Py_INCREF(elem);
2920 PyTuple_SET_ITEM(result, i, elem);
2921 }
2922 Py_DECREF(old_result);
2923 }
2924 /* Now, we've got the only copy so we can update it in-place CPython's
2925 empty tuple is a singleton and cached in PyTuple's freelist. */
2926 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002927
Tim Peters9edb1682013-09-03 11:49:31 -05002928 /* Scan indices right-to-left until finding one that is not
2929 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002930 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2931 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002932
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002934 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002935 if (i < 0)
2936 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002937
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002938 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002939 maximum. Then set all to the right to the same value. */
2940 index = indices[i] + 1;
2941 assert(index < n);
2942 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002943 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002944 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002945 Py_INCREF(elem);
2946 oldelem = PyTuple_GET_ITEM(result, i);
2947 PyTuple_SET_ITEM(result, i, elem);
2948 Py_DECREF(oldelem);
2949 }
2950 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002951
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002952 Py_INCREF(result);
2953 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002954
2955empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002956 co->stopped = 1;
2957 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002958}
2959
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002960static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302961cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002962{
2963 if (lz->result == NULL) {
2964 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2965 } else if (lz->stopped) {
2966 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2967 } else {
2968 PyObject *indices;
2969 Py_ssize_t i;
2970
2971 /* we must pickle the indices and use them for setstate */
2972 indices = PyTuple_New(lz->r);
2973 if (!indices)
2974 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002975 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002976 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2977 if (!index) {
2978 Py_DECREF(indices);
2979 return NULL;
2980 }
2981 PyTuple_SET_ITEM(indices, i, index);
2982 }
2983
2984 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2985 }
2986}
2987
2988static PyObject *
2989cwr_setstate(cwrobject *lz, PyObject *state)
2990{
2991 PyObject *result;
2992 Py_ssize_t n, i;
2993
2994 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2995 {
2996 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2997 return NULL;
2998 }
2999
3000 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003001 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003002 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
3003 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003004
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003005 if (index < 0 && PyErr_Occurred())
3006 return NULL; /* not an integer */
3007 /* clamp the index */
3008 if (index < 0)
3009 index = 0;
3010 else if (index > n-1)
3011 index = n-1;
3012 lz->indices[i] = index;
3013 }
3014 result = PyTuple_New(lz->r);
3015 if (result == NULL)
3016 return NULL;
3017 for (i=0; i<lz->r; i++) {
3018 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3019 Py_INCREF(element);
3020 PyTuple_SET_ITEM(result, i, element);
3021 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003022 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003023 Py_RETURN_NONE;
3024}
3025
3026static PyMethodDef cwr_methods[] = {
3027 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3028 reduce_doc},
3029 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3030 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003031 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3032 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003033 {NULL, NULL} /* sentinel */
3034};
3035
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003036static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003037 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003038 "itertools.combinations_with_replacement", /* tp_name */
3039 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003040 0, /* tp_itemsize */
3041 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003042 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003043 0, /* tp_print */
3044 0, /* tp_getattr */
3045 0, /* tp_setattr */
3046 0, /* tp_reserved */
3047 0, /* tp_repr */
3048 0, /* tp_as_number */
3049 0, /* tp_as_sequence */
3050 0, /* tp_as_mapping */
3051 0, /* tp_hash */
3052 0, /* tp_call */
3053 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003054 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003055 0, /* tp_setattro */
3056 0, /* tp_as_buffer */
3057 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003058 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003059 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003060 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 0, /* tp_clear */
3062 0, /* tp_richcompare */
3063 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003064 PyObject_SelfIter, /* tp_iter */
3065 (iternextfunc)cwr_next, /* tp_iternext */
3066 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003067 0, /* tp_members */
3068 0, /* tp_getset */
3069 0, /* tp_base */
3070 0, /* tp_dict */
3071 0, /* tp_descr_get */
3072 0, /* tp_descr_set */
3073 0, /* tp_dictoffset */
3074 0, /* tp_init */
3075 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003076 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003077 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003078};
3079
3080
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003081/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003082
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003083def permutations(iterable, r=None):
3084 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3085 pool = tuple(iterable)
3086 n = len(pool)
3087 r = n if r is None else r
3088 indices = range(n)
3089 cycles = range(n-r+1, n+1)[::-1]
3090 yield tuple(pool[i] for i in indices[:r])
3091 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003092 for i in reversed(range(r)):
3093 cycles[i] -= 1
3094 if cycles[i] == 0:
3095 indices[i:] = indices[i+1:] + indices[i:i+1]
3096 cycles[i] = n - i
3097 else:
3098 j = cycles[i]
3099 indices[i], indices[-j] = indices[-j], indices[i]
3100 yield tuple(pool[i] for i in indices[:r])
3101 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003102 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003103 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003104*/
3105
3106typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003107 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003108 PyObject *pool; /* input converted to a tuple */
3109 Py_ssize_t *indices; /* one index per element in the pool */
3110 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3111 PyObject *result; /* most recently returned result tuple */
3112 Py_ssize_t r; /* size of result tuple */
3113 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003114} permutationsobject;
3115
3116static PyTypeObject permutations_type;
3117
Tal Einatc4bccd32018-09-12 00:49:13 +03003118/*[clinic input]
3119@classmethod
3120itertools.permutations.__new__
3121 iterable: object
3122 r as robj: object = None
3123Return successive r-length permutations of elements in the iterable.
3124
3125permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3126[clinic start generated code]*/
3127
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003128static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003129itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3130 PyObject *robj)
3131/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003132{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003133 permutationsobject *po;
3134 Py_ssize_t n;
3135 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003136 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003137 Py_ssize_t *indices = NULL;
3138 Py_ssize_t *cycles = NULL;
3139 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003141 pool = PySequence_Tuple(iterable);
3142 if (pool == NULL)
3143 goto error;
3144 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003146 r = n;
3147 if (robj != Py_None) {
3148 if (!PyLong_Check(robj)) {
3149 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3150 goto error;
3151 }
3152 r = PyLong_AsSsize_t(robj);
3153 if (r == -1 && PyErr_Occurred())
3154 goto error;
3155 }
3156 if (r < 0) {
3157 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3158 goto error;
3159 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003160
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003161 indices = PyMem_New(Py_ssize_t, n);
3162 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003163 if (indices == NULL || cycles == NULL) {
3164 PyErr_NoMemory();
3165 goto error;
3166 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003167
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003168 for (i=0 ; i<n ; i++)
3169 indices[i] = i;
3170 for (i=0 ; i<r ; i++)
3171 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003172
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003173 /* create permutationsobject structure */
3174 po = (permutationsobject *)type->tp_alloc(type, 0);
3175 if (po == NULL)
3176 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003178 po->pool = pool;
3179 po->indices = indices;
3180 po->cycles = cycles;
3181 po->result = NULL;
3182 po->r = r;
3183 po->stopped = r > n ? 1 : 0;
3184
3185 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003186
3187error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003188 if (indices != NULL)
3189 PyMem_Free(indices);
3190 if (cycles != NULL)
3191 PyMem_Free(cycles);
3192 Py_XDECREF(pool);
3193 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003194}
3195
3196static void
3197permutations_dealloc(permutationsobject *po)
3198{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003199 PyObject_GC_UnTrack(po);
3200 Py_XDECREF(po->pool);
3201 Py_XDECREF(po->result);
3202 PyMem_Free(po->indices);
3203 PyMem_Free(po->cycles);
3204 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003205}
3206
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003207static PyObject *
3208permutations_sizeof(permutationsobject *po, void *unused)
3209{
3210 Py_ssize_t res;
3211
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003212 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003213 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3214 res += po->r * sizeof(Py_ssize_t);
3215 return PyLong_FromSsize_t(res);
3216}
3217
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003218static int
3219permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3220{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003221 Py_VISIT(po->pool);
3222 Py_VISIT(po->result);
3223 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003224}
3225
3226static PyObject *
3227permutations_next(permutationsobject *po)
3228{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003229 PyObject *elem;
3230 PyObject *oldelem;
3231 PyObject *pool = po->pool;
3232 Py_ssize_t *indices = po->indices;
3233 Py_ssize_t *cycles = po->cycles;
3234 PyObject *result = po->result;
3235 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3236 Py_ssize_t r = po->r;
3237 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003238
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003239 if (po->stopped)
3240 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003242 if (result == NULL) {
3243 /* On the first pass, initialize result tuple using the indices */
3244 result = PyTuple_New(r);
3245 if (result == NULL)
3246 goto empty;
3247 po->result = result;
3248 for (i=0; i<r ; i++) {
3249 index = indices[i];
3250 elem = PyTuple_GET_ITEM(pool, index);
3251 Py_INCREF(elem);
3252 PyTuple_SET_ITEM(result, i, elem);
3253 }
3254 } else {
3255 if (n == 0)
3256 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003257
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003258 /* Copy the previous result tuple or re-use it if available */
3259 if (Py_REFCNT(result) > 1) {
3260 PyObject *old_result = result;
3261 result = PyTuple_New(r);
3262 if (result == NULL)
3263 goto empty;
3264 po->result = result;
3265 for (i=0; i<r ; i++) {
3266 elem = PyTuple_GET_ITEM(old_result, i);
3267 Py_INCREF(elem);
3268 PyTuple_SET_ITEM(result, i, elem);
3269 }
3270 Py_DECREF(old_result);
3271 }
3272 /* Now, we've got the only copy so we can update it in-place */
3273 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003274
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003275 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3276 for (i=r-1 ; i>=0 ; i--) {
3277 cycles[i] -= 1;
3278 if (cycles[i] == 0) {
3279 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3280 index = indices[i];
3281 for (j=i ; j<n-1 ; j++)
3282 indices[j] = indices[j+1];
3283 indices[n-1] = index;
3284 cycles[i] = n - i;
3285 } else {
3286 j = cycles[i];
3287 index = indices[i];
3288 indices[i] = indices[n-j];
3289 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003290
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003291 for (k=i; k<r ; k++) {
3292 /* start with i, the leftmost element that changed */
3293 /* yield tuple(pool[k] for k in indices[:r]) */
3294 index = indices[k];
3295 elem = PyTuple_GET_ITEM(pool, index);
3296 Py_INCREF(elem);
3297 oldelem = PyTuple_GET_ITEM(result, k);
3298 PyTuple_SET_ITEM(result, k, elem);
3299 Py_DECREF(oldelem);
3300 }
3301 break;
3302 }
3303 }
3304 /* If i is negative, then the cycles have all
3305 rolled-over and we're done. */
3306 if (i < 0)
3307 goto empty;
3308 }
3309 Py_INCREF(result);
3310 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003311
3312empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003313 po->stopped = 1;
3314 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003315}
3316
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003317static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303318permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003319{
3320 if (po->result == NULL) {
3321 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3322 } else if (po->stopped) {
3323 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3324 } else {
3325 PyObject *indices=NULL, *cycles=NULL;
3326 Py_ssize_t n, i;
3327
3328 /* we must pickle the indices and cycles and use them for setstate */
3329 n = PyTuple_GET_SIZE(po->pool);
3330 indices = PyTuple_New(n);
3331 if (indices == NULL)
3332 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003333 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003334 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3335 if (!index)
3336 goto err;
3337 PyTuple_SET_ITEM(indices, i, index);
3338 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003339
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003340 cycles = PyTuple_New(po->r);
3341 if (cycles == NULL)
3342 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003343 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003344 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3345 if (!index)
3346 goto err;
3347 PyTuple_SET_ITEM(cycles, i, index);
3348 }
3349 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3350 po->pool, po->r,
3351 indices, cycles);
3352 err:
3353 Py_XDECREF(indices);
3354 Py_XDECREF(cycles);
3355 return NULL;
3356 }
3357}
3358
3359static PyObject *
3360permutations_setstate(permutationsobject *po, PyObject *state)
3361{
3362 PyObject *indices, *cycles, *result;
3363 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003364
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003365 if (!PyTuple_Check(state)) {
3366 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3367 return NULL;
3368 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003369 if (!PyArg_ParseTuple(state, "O!O!",
3370 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003371 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003372 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003373 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003374
3375 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003376 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003377 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3378 return NULL;
3379 }
3380
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003381 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003382 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3383 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3384 if (index < 0 && PyErr_Occurred())
3385 return NULL; /* not an integer */
3386 /* clamp the index */
3387 if (index < 0)
3388 index = 0;
3389 else if (index > n-1)
3390 index = n-1;
3391 po->indices[i] = index;
3392 }
3393
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003394 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003395 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3396 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3397 if (index < 0 && PyErr_Occurred())
3398 return NULL; /* not an integer */
3399 if (index < 1)
3400 index = 1;
3401 else if (index > n-i)
3402 index = n-i;
3403 po->cycles[i] = index;
3404 }
3405 result = PyTuple_New(po->r);
3406 if (result == NULL)
3407 return NULL;
3408 for (i=0; i<po->r; i++) {
3409 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3410 Py_INCREF(element);
3411 PyTuple_SET_ITEM(result, i, element);
3412 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003413 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003414 Py_RETURN_NONE;
3415}
3416
3417static PyMethodDef permuations_methods[] = {
3418 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3419 reduce_doc},
3420 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3421 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003422 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3423 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003424 {NULL, NULL} /* sentinel */
3425};
3426
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003427static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003428 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003429 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003430 sizeof(permutationsobject), /* tp_basicsize */
3431 0, /* tp_itemsize */
3432 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003433 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003434 0, /* tp_print */
3435 0, /* tp_getattr */
3436 0, /* tp_setattr */
3437 0, /* tp_reserved */
3438 0, /* tp_repr */
3439 0, /* tp_as_number */
3440 0, /* tp_as_sequence */
3441 0, /* tp_as_mapping */
3442 0, /* tp_hash */
3443 0, /* tp_call */
3444 0, /* tp_str */
3445 PyObject_GenericGetAttr, /* tp_getattro */
3446 0, /* tp_setattro */
3447 0, /* tp_as_buffer */
3448 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3449 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003450 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003451 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003452 0, /* tp_clear */
3453 0, /* tp_richcompare */
3454 0, /* tp_weaklistoffset */
3455 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003456 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003457 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003458 0, /* tp_members */
3459 0, /* tp_getset */
3460 0, /* tp_base */
3461 0, /* tp_dict */
3462 0, /* tp_descr_get */
3463 0, /* tp_descr_set */
3464 0, /* tp_dictoffset */
3465 0, /* tp_init */
3466 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003467 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003468 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003469};
3470
Tal Einatc4bccd32018-09-12 00:49:13 +03003471
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003472/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003473
3474typedef struct {
3475 PyObject_HEAD
3476 PyObject *total;
3477 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003478 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003479 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003480} accumulateobject;
3481
3482static PyTypeObject accumulate_type;
3483
Tal Einatc4bccd32018-09-12 00:49:13 +03003484/*[clinic input]
3485@classmethod
3486itertools.accumulate.__new__
3487 iterable: object
3488 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003489 *
3490 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003491Return series of accumulated sums (or other binary function results).
3492[clinic start generated code]*/
3493
Raymond Hettinger482ba772010-12-01 22:48:00 +00003494static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003495itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003496 PyObject *binop, PyObject *initial)
3497/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003498{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003499 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003500 accumulateobject *lz;
3501
Raymond Hettinger482ba772010-12-01 22:48:00 +00003502 /* Get iterator. */
3503 it = PyObject_GetIter(iterable);
3504 if (it == NULL)
3505 return NULL;
3506
Raymond Hettinger482ba772010-12-01 22:48:00 +00003507 /* create accumulateobject structure */
3508 lz = (accumulateobject *)type->tp_alloc(type, 0);
3509 if (lz == NULL) {
3510 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003511 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003512 }
3513
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003514 if (binop != Py_None) {
3515 Py_XINCREF(binop);
3516 lz->binop = binop;
3517 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003518 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003519 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003520 Py_XINCREF(initial);
3521 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003522 return (PyObject *)lz;
3523}
3524
3525static void
3526accumulate_dealloc(accumulateobject *lz)
3527{
3528 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003529 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003530 Py_XDECREF(lz->total);
3531 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003532 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003533 Py_TYPE(lz)->tp_free(lz);
3534}
3535
3536static int
3537accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3538{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003539 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003540 Py_VISIT(lz->it);
3541 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003542 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003543 return 0;
3544}
3545
3546static PyObject *
3547accumulate_next(accumulateobject *lz)
3548{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003549 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003550
Lisa Roach9718b592018-09-23 17:34:59 -07003551 if (lz->initial != Py_None) {
3552 lz->total = lz->initial;
3553 Py_INCREF(Py_None);
3554 lz->initial = Py_None;
3555 Py_INCREF(lz->total);
3556 return lz->total;
3557 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003558 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003559 if (val == NULL)
3560 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003561
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003562 if (lz->total == NULL) {
3563 Py_INCREF(val);
3564 lz->total = val;
3565 return lz->total;
3566 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003567
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003568 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003569 newtotal = PyNumber_Add(lz->total, val);
3570 else
3571 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003572 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003573 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003574 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003575
Raymond Hettinger482ba772010-12-01 22:48:00 +00003576 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003577 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003578 return newtotal;
3579}
3580
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003581static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303582accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003583{
Lisa Roach9718b592018-09-23 17:34:59 -07003584 if (lz->initial != Py_None) {
3585 PyObject *it;
3586
3587 assert(lz->total == NULL);
3588 if (PyType_Ready(&chain_type) < 0)
3589 return NULL;
3590 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3591 lz->initial, lz->it);
3592 if (it == NULL)
3593 return NULL;
3594 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3595 it, lz->binop?lz->binop:Py_None, Py_None);
3596 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003597 if (lz->total == Py_None) {
3598 PyObject *it;
3599
3600 if (PyType_Ready(&chain_type) < 0)
3601 return NULL;
3602 if (PyType_Ready(&islice_type) < 0)
3603 return NULL;
3604 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3605 lz->total, lz->it);
3606 if (it == NULL)
3607 return NULL;
3608 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3609 it, lz->binop ? lz->binop : Py_None);
3610 if (it == NULL)
3611 return NULL;
3612 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3613 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003614 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3615 lz->it, lz->binop?lz->binop:Py_None,
3616 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003617}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003618
3619static PyObject *
3620accumulate_setstate(accumulateobject *lz, PyObject *state)
3621{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003622 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003623 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003624 Py_RETURN_NONE;
3625}
3626
3627static PyMethodDef accumulate_methods[] = {
3628 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3629 reduce_doc},
3630 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3631 setstate_doc},
3632 {NULL, NULL} /* sentinel */
3633};
3634
Raymond Hettinger482ba772010-12-01 22:48:00 +00003635static PyTypeObject accumulate_type = {
3636 PyVarObject_HEAD_INIT(NULL, 0)
3637 "itertools.accumulate", /* tp_name */
3638 sizeof(accumulateobject), /* tp_basicsize */
3639 0, /* tp_itemsize */
3640 /* methods */
3641 (destructor)accumulate_dealloc, /* tp_dealloc */
3642 0, /* tp_print */
3643 0, /* tp_getattr */
3644 0, /* tp_setattr */
3645 0, /* tp_reserved */
3646 0, /* tp_repr */
3647 0, /* tp_as_number */
3648 0, /* tp_as_sequence */
3649 0, /* tp_as_mapping */
3650 0, /* tp_hash */
3651 0, /* tp_call */
3652 0, /* tp_str */
3653 PyObject_GenericGetAttr, /* tp_getattro */
3654 0, /* tp_setattro */
3655 0, /* tp_as_buffer */
3656 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3657 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003658 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003659 (traverseproc)accumulate_traverse, /* tp_traverse */
3660 0, /* tp_clear */
3661 0, /* tp_richcompare */
3662 0, /* tp_weaklistoffset */
3663 PyObject_SelfIter, /* tp_iter */
3664 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003665 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003666 0, /* tp_members */
3667 0, /* tp_getset */
3668 0, /* tp_base */
3669 0, /* tp_dict */
3670 0, /* tp_descr_get */
3671 0, /* tp_descr_set */
3672 0, /* tp_dictoffset */
3673 0, /* tp_init */
3674 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003675 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003676 PyObject_GC_Del, /* tp_free */
3677};
3678
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003679
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003680/* compress object ************************************************************/
3681
3682/* Equivalent to:
3683
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003684 def compress(data, selectors):
3685 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3686 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003687*/
3688
3689typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003690 PyObject_HEAD
3691 PyObject *data;
3692 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003693} compressobject;
3694
3695static PyTypeObject compress_type;
3696
Tal Einatc4bccd32018-09-12 00:49:13 +03003697/*[clinic input]
3698@classmethod
3699itertools.compress.__new__
3700 data as seq1: object
3701 selectors as seq2: object
3702Return data elements corresponding to true selector elements.
3703
3704Forms a shorter iterator from selected data elements using the selectors to
3705choose the data elements.
3706[clinic start generated code]*/
3707
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003708static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003709itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3710/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003711{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003712 PyObject *data=NULL, *selectors=NULL;
3713 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003714
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003715 data = PyObject_GetIter(seq1);
3716 if (data == NULL)
3717 goto fail;
3718 selectors = PyObject_GetIter(seq2);
3719 if (selectors == NULL)
3720 goto fail;
3721
3722 /* create compressobject structure */
3723 lz = (compressobject *)type->tp_alloc(type, 0);
3724 if (lz == NULL)
3725 goto fail;
3726 lz->data = data;
3727 lz->selectors = selectors;
3728 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003729
3730fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003731 Py_XDECREF(data);
3732 Py_XDECREF(selectors);
3733 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003734}
3735
3736static void
3737compress_dealloc(compressobject *lz)
3738{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003739 PyObject_GC_UnTrack(lz);
3740 Py_XDECREF(lz->data);
3741 Py_XDECREF(lz->selectors);
3742 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003743}
3744
3745static int
3746compress_traverse(compressobject *lz, visitproc visit, void *arg)
3747{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003748 Py_VISIT(lz->data);
3749 Py_VISIT(lz->selectors);
3750 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003751}
3752
3753static PyObject *
3754compress_next(compressobject *lz)
3755{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003756 PyObject *data = lz->data, *selectors = lz->selectors;
3757 PyObject *datum, *selector;
3758 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3759 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3760 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003761
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003762 while (1) {
3763 /* Steps: get datum, get selector, evaluate selector.
3764 Order is important (to match the pure python version
3765 in terms of which input gets a chance to raise an
3766 exception first).
3767 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003768
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003769 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003770 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003771 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003772
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003773 selector = selectornext(selectors);
3774 if (selector == NULL) {
3775 Py_DECREF(datum);
3776 return NULL;
3777 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003778
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003779 ok = PyObject_IsTrue(selector);
3780 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003781 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003782 return datum;
3783 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003784 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003785 return NULL;
3786 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003787}
3788
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003789static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303790compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003791{
3792 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3793 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003794}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003795
3796static PyMethodDef compress_methods[] = {
3797 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3798 reduce_doc},
3799 {NULL, NULL} /* sentinel */
3800};
3801
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003802static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003803 PyVarObject_HEAD_INIT(NULL, 0)
3804 "itertools.compress", /* tp_name */
3805 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003806 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003807 /* methods */
3808 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003809 0, /* tp_print */
3810 0, /* tp_getattr */
3811 0, /* tp_setattr */
3812 0, /* tp_reserved */
3813 0, /* tp_repr */
3814 0, /* tp_as_number */
3815 0, /* tp_as_sequence */
3816 0, /* tp_as_mapping */
3817 0, /* tp_hash */
3818 0, /* tp_call */
3819 0, /* tp_str */
3820 PyObject_GenericGetAttr, /* tp_getattro */
3821 0, /* tp_setattro */
3822 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003823 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003824 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003825 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003826 (traverseproc)compress_traverse, /* tp_traverse */
3827 0, /* tp_clear */
3828 0, /* tp_richcompare */
3829 0, /* tp_weaklistoffset */
3830 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003831 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003832 compress_methods, /* tp_methods */
3833 0, /* tp_members */
3834 0, /* tp_getset */
3835 0, /* tp_base */
3836 0, /* tp_dict */
3837 0, /* tp_descr_get */
3838 0, /* tp_descr_set */
3839 0, /* tp_dictoffset */
3840 0, /* tp_init */
3841 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003842 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003843 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003844};
3845
3846
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003847/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003848
3849typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003850 PyObject_HEAD
3851 PyObject *func;
3852 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003853} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003854
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003855static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003856
Tal Einatc4bccd32018-09-12 00:49:13 +03003857/*[clinic input]
3858@classmethod
3859itertools.filterfalse.__new__
3860 function as func: object
3861 iterable as seq: object
3862 /
3863Return those items of iterable for which function(item) is false.
3864
3865If function is None, return the items that are false.
3866[clinic start generated code]*/
3867
Raymond Hettinger60eca932003-02-09 06:40:58 +00003868static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003869itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3870/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003871{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003872 PyObject *it;
3873 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003875 /* Get iterator. */
3876 it = PyObject_GetIter(seq);
3877 if (it == NULL)
3878 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003880 /* create filterfalseobject structure */
3881 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3882 if (lz == NULL) {
3883 Py_DECREF(it);
3884 return NULL;
3885 }
3886 Py_INCREF(func);
3887 lz->func = func;
3888 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003889
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003890 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003891}
3892
3893static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003894filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003895{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003896 PyObject_GC_UnTrack(lz);
3897 Py_XDECREF(lz->func);
3898 Py_XDECREF(lz->it);
3899 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003900}
3901
3902static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003903filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003904{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003905 Py_VISIT(lz->it);
3906 Py_VISIT(lz->func);
3907 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003908}
3909
3910static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003911filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 PyObject *item;
3914 PyObject *it = lz->it;
3915 long ok;
3916 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003917
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003918 iternext = *Py_TYPE(it)->tp_iternext;
3919 for (;;) {
3920 item = iternext(it);
3921 if (item == NULL)
3922 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003923
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003924 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3925 ok = PyObject_IsTrue(item);
3926 } else {
3927 PyObject *good;
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01003928 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003929 if (good == NULL) {
3930 Py_DECREF(item);
3931 return NULL;
3932 }
3933 ok = PyObject_IsTrue(good);
3934 Py_DECREF(good);
3935 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003936 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 return item;
3938 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003939 if (ok < 0)
3940 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003941 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003942}
3943
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003944static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303945filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003946{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003947 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3948}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003949
3950static PyMethodDef filterfalse_methods[] = {
3951 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3952 reduce_doc},
3953 {NULL, NULL} /* sentinel */
3954};
3955
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003956static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 PyVarObject_HEAD_INIT(NULL, 0)
3958 "itertools.filterfalse", /* tp_name */
3959 sizeof(filterfalseobject), /* tp_basicsize */
3960 0, /* tp_itemsize */
3961 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003962 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003963 0, /* tp_print */
3964 0, /* tp_getattr */
3965 0, /* tp_setattr */
3966 0, /* tp_reserved */
3967 0, /* tp_repr */
3968 0, /* tp_as_number */
3969 0, /* tp_as_sequence */
3970 0, /* tp_as_mapping */
3971 0, /* tp_hash */
3972 0, /* tp_call */
3973 0, /* tp_str */
3974 PyObject_GenericGetAttr, /* tp_getattro */
3975 0, /* tp_setattro */
3976 0, /* tp_as_buffer */
3977 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3978 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003979 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003980 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003981 0, /* tp_clear */
3982 0, /* tp_richcompare */
3983 0, /* tp_weaklistoffset */
3984 PyObject_SelfIter, /* tp_iter */
3985 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003986 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003987 0, /* tp_members */
3988 0, /* tp_getset */
3989 0, /* tp_base */
3990 0, /* tp_dict */
3991 0, /* tp_descr_get */
3992 0, /* tp_descr_set */
3993 0, /* tp_dictoffset */
3994 0, /* tp_init */
3995 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003996 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003997 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003998};
3999
4000
4001/* count object ************************************************************/
4002
4003typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004004 PyObject_HEAD
4005 Py_ssize_t cnt;
4006 PyObject *long_cnt;
4007 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004008} countobject;
4009
Raymond Hettinger30729212009-02-12 06:28:27 +00004010/* Counting logic and invariants:
4011
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004012fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00004013
Serhiy Storchaka483405b2015-02-17 10:14:30 +02004014 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004015 Advances with: cnt += 1
4016 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00004017
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004018slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004019
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004020 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4021 All counting is done with python objects (no overflows or underflows).
4022 Advances with: long_cnt += long_step
4023 Step may be zero -- effectively a slow version of repeat(cnt).
4024 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004025*/
4026
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004027static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004028
Tal Einatc4bccd32018-09-12 00:49:13 +03004029/*[clinic input]
4030@classmethod
4031itertools.count.__new__
4032 start as long_cnt: object(c_default="NULL") = 0
4033 step as long_step: object(c_default="NULL") = 1
4034Return a count object whose .__next__() method returns consecutive values.
4035
4036Equivalent to:
4037 def count(firstval=0, step=1):
4038 x = firstval
4039 while 1:
4040 yield x
4041 x += step
4042[clinic start generated code]*/
4043
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004044static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004045itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4046 PyObject *long_step)
4047/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004048{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004049 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004050 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004051 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004052 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004054 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4055 (long_step != NULL && !PyNumber_Check(long_step))) {
4056 PyErr_SetString(PyExc_TypeError, "a number is required");
4057 return NULL;
4058 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004059
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004060 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4061 (long_step == NULL || PyLong_Check(long_step));
4062
4063 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004064 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004065 if (fast_mode) {
4066 assert(PyLong_Check(long_cnt));
4067 cnt = PyLong_AsSsize_t(long_cnt);
4068 if (cnt == -1 && PyErr_Occurred()) {
4069 PyErr_Clear();
4070 fast_mode = 0;
4071 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004072 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004073 } else {
4074 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004075 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004076 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004077 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004079 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004080 if (long_step == NULL)
4081 long_step = _PyLong_One;
4082 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004083
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004084 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004085
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004086 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004087 if (fast_mode) {
4088 assert(PyLong_Check(long_step));
4089 step = PyLong_AsLong(long_step);
4090 if (step != 1) {
4091 fast_mode = 0;
4092 if (step == -1 && PyErr_Occurred())
4093 PyErr_Clear();
4094 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004095 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004096
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004097 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004098 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004099 else
4100 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004101
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004102 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4103 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4104 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004106
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004107 /* create countobject structure */
4108 lz = (countobject *)type->tp_alloc(type, 0);
4109 if (lz == NULL) {
4110 Py_XDECREF(long_cnt);
4111 return NULL;
4112 }
4113 lz->cnt = cnt;
4114 lz->long_cnt = long_cnt;
4115 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004117 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004118}
4119
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004120static void
4121count_dealloc(countobject *lz)
4122{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 PyObject_GC_UnTrack(lz);
4124 Py_XDECREF(lz->long_cnt);
4125 Py_XDECREF(lz->long_step);
4126 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004127}
4128
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004129static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004130count_traverse(countobject *lz, visitproc visit, void *arg)
4131{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004132 Py_VISIT(lz->long_cnt);
4133 Py_VISIT(lz->long_step);
4134 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004135}
4136
4137static PyObject *
4138count_nextlong(countobject *lz)
4139{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004140 PyObject *long_cnt;
4141 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004142
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004143 long_cnt = lz->long_cnt;
4144 if (long_cnt == NULL) {
4145 /* Switch to slow_mode */
4146 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4147 if (long_cnt == NULL)
4148 return NULL;
4149 }
4150 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004151
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004152 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4153 if (stepped_up == NULL)
4154 return NULL;
4155 lz->long_cnt = stepped_up;
4156 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004157}
4158
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004159static PyObject *
4160count_next(countobject *lz)
4161{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004162 if (lz->cnt == PY_SSIZE_T_MAX)
4163 return count_nextlong(lz);
4164 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004165}
4166
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004167static PyObject *
4168count_repr(countobject *lz)
4169{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004170 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004171 return PyUnicode_FromFormat("%s(%zd)",
4172 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004173
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004174 if (PyLong_Check(lz->long_step)) {
4175 long step = PyLong_AsLong(lz->long_step);
4176 if (step == -1 && PyErr_Occurred()) {
4177 PyErr_Clear();
4178 }
4179 if (step == 1) {
4180 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004181 return PyUnicode_FromFormat("%s(%R)",
4182 _PyType_Name(Py_TYPE(lz)),
4183 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004184 }
4185 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004186 return PyUnicode_FromFormat("%s(%R, %R)",
4187 _PyType_Name(Py_TYPE(lz)),
4188 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004189}
4190
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004191static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304192count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004193{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004194 if (lz->cnt == PY_SSIZE_T_MAX)
4195 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4196 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004197}
4198
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004199static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004200 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004201 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004202 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004203};
4204
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004205static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004206 PyVarObject_HEAD_INIT(NULL, 0)
4207 "itertools.count", /* tp_name */
4208 sizeof(countobject), /* tp_basicsize */
4209 0, /* tp_itemsize */
4210 /* methods */
4211 (destructor)count_dealloc, /* tp_dealloc */
4212 0, /* tp_print */
4213 0, /* tp_getattr */
4214 0, /* tp_setattr */
4215 0, /* tp_reserved */
4216 (reprfunc)count_repr, /* tp_repr */
4217 0, /* tp_as_number */
4218 0, /* tp_as_sequence */
4219 0, /* tp_as_mapping */
4220 0, /* tp_hash */
4221 0, /* tp_call */
4222 0, /* tp_str */
4223 PyObject_GenericGetAttr, /* tp_getattro */
4224 0, /* tp_setattro */
4225 0, /* tp_as_buffer */
4226 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004227 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004228 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004229 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004230 0, /* tp_clear */
4231 0, /* tp_richcompare */
4232 0, /* tp_weaklistoffset */
4233 PyObject_SelfIter, /* tp_iter */
4234 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004235 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004236 0, /* tp_members */
4237 0, /* tp_getset */
4238 0, /* tp_base */
4239 0, /* tp_dict */
4240 0, /* tp_descr_get */
4241 0, /* tp_descr_set */
4242 0, /* tp_dictoffset */
4243 0, /* tp_init */
4244 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004245 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004246 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004247};
4248
4249
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004250/* repeat object ************************************************************/
4251
4252typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004253 PyObject_HEAD
4254 PyObject *element;
4255 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004256} repeatobject;
4257
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004258static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004259
4260static PyObject *
4261repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4262{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004263 repeatobject *ro;
4264 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004265 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004266 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004267
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004268 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4269 &element, &cnt))
4270 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004271
Raymond Hettinger97d35552014-06-24 21:36:58 -07004272 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004273 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004274 /* Does user supply times argument? */
4275 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004276 cnt = 0;
4277
4278 ro = (repeatobject *)type->tp_alloc(type, 0);
4279 if (ro == NULL)
4280 return NULL;
4281 Py_INCREF(element);
4282 ro->element = element;
4283 ro->cnt = cnt;
4284 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004285}
4286
4287static void
4288repeat_dealloc(repeatobject *ro)
4289{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004290 PyObject_GC_UnTrack(ro);
4291 Py_XDECREF(ro->element);
4292 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004293}
4294
4295static int
4296repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4297{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004298 Py_VISIT(ro->element);
4299 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004300}
4301
4302static PyObject *
4303repeat_next(repeatobject *ro)
4304{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004305 if (ro->cnt == 0)
4306 return NULL;
4307 if (ro->cnt > 0)
4308 ro->cnt--;
4309 Py_INCREF(ro->element);
4310 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004311}
4312
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004313static PyObject *
4314repeat_repr(repeatobject *ro)
4315{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004316 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004317 return PyUnicode_FromFormat("%s(%R)",
4318 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004319 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004320 return PyUnicode_FromFormat("%s(%R, %zd)",
4321 _PyType_Name(Py_TYPE(ro)), ro->element,
4322 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004323}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004324
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004325static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304326repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004327{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004328 if (ro->cnt == -1) {
4329 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4330 return NULL;
4331 }
4332 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004333}
4334
Armin Rigof5b3e362006-02-11 21:32:43 +00004335PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004336
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004337static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304338repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004339{
4340 /* unpickle this so that a new repeat iterator is constructed with an
4341 * object, then call __setstate__ on it to set cnt
4342 */
4343 if (ro->cnt >= 0)
4344 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4345 else
4346 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4347}
4348
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004349static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004350 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004351 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004352 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004353};
4354
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004355PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004356"repeat(object [,times]) -> create an iterator which returns the object\n\
4357for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004358endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004359
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004360static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004361 PyVarObject_HEAD_INIT(NULL, 0)
4362 "itertools.repeat", /* tp_name */
4363 sizeof(repeatobject), /* tp_basicsize */
4364 0, /* tp_itemsize */
4365 /* methods */
4366 (destructor)repeat_dealloc, /* tp_dealloc */
4367 0, /* tp_print */
4368 0, /* tp_getattr */
4369 0, /* tp_setattr */
4370 0, /* tp_reserved */
4371 (reprfunc)repeat_repr, /* tp_repr */
4372 0, /* tp_as_number */
4373 0, /* tp_as_sequence */
4374 0, /* tp_as_mapping */
4375 0, /* tp_hash */
4376 0, /* tp_call */
4377 0, /* tp_str */
4378 PyObject_GenericGetAttr, /* tp_getattro */
4379 0, /* tp_setattro */
4380 0, /* tp_as_buffer */
4381 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4382 Py_TPFLAGS_BASETYPE, /* tp_flags */
4383 repeat_doc, /* tp_doc */
4384 (traverseproc)repeat_traverse, /* tp_traverse */
4385 0, /* tp_clear */
4386 0, /* tp_richcompare */
4387 0, /* tp_weaklistoffset */
4388 PyObject_SelfIter, /* tp_iter */
4389 (iternextfunc)repeat_next, /* tp_iternext */
4390 repeat_methods, /* tp_methods */
4391 0, /* tp_members */
4392 0, /* tp_getset */
4393 0, /* tp_base */
4394 0, /* tp_dict */
4395 0, /* tp_descr_get */
4396 0, /* tp_descr_set */
4397 0, /* tp_dictoffset */
4398 0, /* tp_init */
4399 0, /* tp_alloc */
4400 repeat_new, /* tp_new */
4401 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004402};
4403
Tal Einatc4bccd32018-09-12 00:49:13 +03004404
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004405/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004406
4407typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004408 PyObject_HEAD
4409 Py_ssize_t tuplesize;
4410 Py_ssize_t numactive;
4411 PyObject *ittuple; /* tuple of iterators */
4412 PyObject *result;
4413 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004414} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004415
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004416static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004417
4418static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004419zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004420{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004421 ziplongestobject *lz;
4422 Py_ssize_t i;
4423 PyObject *ittuple; /* tuple of iterators */
4424 PyObject *result;
4425 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004426 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004427
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004428 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004429 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004430 if (fillvalue == NULL || PyDict_GET_SIZE(kwds) > 1) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004431 PyErr_SetString(PyExc_TypeError,
4432 "zip_longest() got an unexpected keyword argument");
4433 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004434 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004435 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004437 /* args must be a tuple */
4438 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004439 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004440
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004441 /* obtain iterators */
4442 ittuple = PyTuple_New(tuplesize);
4443 if (ittuple == NULL)
4444 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004445 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004446 PyObject *item = PyTuple_GET_ITEM(args, i);
4447 PyObject *it = PyObject_GetIter(item);
4448 if (it == NULL) {
4449 if (PyErr_ExceptionMatches(PyExc_TypeError))
4450 PyErr_Format(PyExc_TypeError,
4451 "zip_longest argument #%zd must support iteration",
4452 i+1);
4453 Py_DECREF(ittuple);
4454 return NULL;
4455 }
4456 PyTuple_SET_ITEM(ittuple, i, it);
4457 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004458
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004459 /* create a result holder */
4460 result = PyTuple_New(tuplesize);
4461 if (result == NULL) {
4462 Py_DECREF(ittuple);
4463 return NULL;
4464 }
4465 for (i=0 ; i < tuplesize ; i++) {
4466 Py_INCREF(Py_None);
4467 PyTuple_SET_ITEM(result, i, Py_None);
4468 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004469
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004470 /* create ziplongestobject structure */
4471 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4472 if (lz == NULL) {
4473 Py_DECREF(ittuple);
4474 Py_DECREF(result);
4475 return NULL;
4476 }
4477 lz->ittuple = ittuple;
4478 lz->tuplesize = tuplesize;
4479 lz->numactive = tuplesize;
4480 lz->result = result;
4481 Py_INCREF(fillvalue);
4482 lz->fillvalue = fillvalue;
4483 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004484}
4485
4486static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004487zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004489 PyObject_GC_UnTrack(lz);
4490 Py_XDECREF(lz->ittuple);
4491 Py_XDECREF(lz->result);
4492 Py_XDECREF(lz->fillvalue);
4493 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004494}
4495
4496static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004497zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004498{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004499 Py_VISIT(lz->ittuple);
4500 Py_VISIT(lz->result);
4501 Py_VISIT(lz->fillvalue);
4502 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004503}
4504
4505static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004506zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004507{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004508 Py_ssize_t i;
4509 Py_ssize_t tuplesize = lz->tuplesize;
4510 PyObject *result = lz->result;
4511 PyObject *it;
4512 PyObject *item;
4513 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004514
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004515 if (tuplesize == 0)
4516 return NULL;
4517 if (lz->numactive == 0)
4518 return NULL;
4519 if (Py_REFCNT(result) == 1) {
4520 Py_INCREF(result);
4521 for (i=0 ; i < tuplesize ; i++) {
4522 it = PyTuple_GET_ITEM(lz->ittuple, i);
4523 if (it == NULL) {
4524 Py_INCREF(lz->fillvalue);
4525 item = lz->fillvalue;
4526 } else {
4527 item = PyIter_Next(it);
4528 if (item == NULL) {
4529 lz->numactive -= 1;
4530 if (lz->numactive == 0 || PyErr_Occurred()) {
4531 lz->numactive = 0;
4532 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004533 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004534 } else {
4535 Py_INCREF(lz->fillvalue);
4536 item = lz->fillvalue;
4537 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4538 Py_DECREF(it);
4539 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004540 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004541 }
4542 olditem = PyTuple_GET_ITEM(result, i);
4543 PyTuple_SET_ITEM(result, i, item);
4544 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004545 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004546 } else {
4547 result = PyTuple_New(tuplesize);
4548 if (result == NULL)
4549 return NULL;
4550 for (i=0 ; i < tuplesize ; i++) {
4551 it = PyTuple_GET_ITEM(lz->ittuple, i);
4552 if (it == NULL) {
4553 Py_INCREF(lz->fillvalue);
4554 item = lz->fillvalue;
4555 } else {
4556 item = PyIter_Next(it);
4557 if (item == NULL) {
4558 lz->numactive -= 1;
4559 if (lz->numactive == 0 || PyErr_Occurred()) {
4560 lz->numactive = 0;
4561 Py_DECREF(result);
4562 return NULL;
4563 } else {
4564 Py_INCREF(lz->fillvalue);
4565 item = lz->fillvalue;
4566 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4567 Py_DECREF(it);
4568 }
4569 }
4570 }
4571 PyTuple_SET_ITEM(result, i, item);
4572 }
4573 }
4574 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004575}
4576
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004577static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304578zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004579{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004580
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004581 /* Create a new tuple with empty sequences where appropriate to pickle.
4582 * Then use setstate to set the fillvalue
4583 */
4584 int i;
4585 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004586
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004587 if (args == NULL)
4588 return NULL;
4589 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4590 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4591 if (elem == NULL) {
4592 elem = PyTuple_New(0);
4593 if (elem == NULL) {
4594 Py_DECREF(args);
4595 return NULL;
4596 }
4597 } else
4598 Py_INCREF(elem);
4599 PyTuple_SET_ITEM(args, i, elem);
4600 }
4601 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4602}
4603
4604static PyObject *
4605zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4606{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004607 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004608 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004609 Py_RETURN_NONE;
4610}
4611
4612static PyMethodDef zip_longest_methods[] = {
4613 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4614 reduce_doc},
4615 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4616 setstate_doc},
4617 {NULL, NULL} /* sentinel */
4618};
4619
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004620PyDoc_STRVAR(zip_longest_doc,
4621"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004622\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004623Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004624the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004625method continues until the longest iterable in the argument sequence\n\
4626is exhausted and then it raises StopIteration. When the shorter iterables\n\
4627are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4628defaults to None or can be specified by a keyword argument.\n\
4629");
4630
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004631static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004632 PyVarObject_HEAD_INIT(NULL, 0)
4633 "itertools.zip_longest", /* tp_name */
4634 sizeof(ziplongestobject), /* tp_basicsize */
4635 0, /* tp_itemsize */
4636 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004637 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004638 0, /* tp_print */
4639 0, /* tp_getattr */
4640 0, /* tp_setattr */
4641 0, /* tp_reserved */
4642 0, /* tp_repr */
4643 0, /* tp_as_number */
4644 0, /* tp_as_sequence */
4645 0, /* tp_as_mapping */
4646 0, /* tp_hash */
4647 0, /* tp_call */
4648 0, /* tp_str */
4649 PyObject_GenericGetAttr, /* tp_getattro */
4650 0, /* tp_setattro */
4651 0, /* tp_as_buffer */
4652 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4653 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004654 zip_longest_doc, /* tp_doc */
4655 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004656 0, /* tp_clear */
4657 0, /* tp_richcompare */
4658 0, /* tp_weaklistoffset */
4659 PyObject_SelfIter, /* tp_iter */
4660 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004661 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004662 0, /* tp_members */
4663 0, /* tp_getset */
4664 0, /* tp_base */
4665 0, /* tp_dict */
4666 0, /* tp_descr_get */
4667 0, /* tp_descr_set */
4668 0, /* tp_dictoffset */
4669 0, /* tp_init */
4670 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004671 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004672 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004673};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004674
Tal Einatc4bccd32018-09-12 00:49:13 +03004675
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004676/* module level code ********************************************************/
4677
4678PyDoc_STRVAR(module_doc,
4679"Functional tools for creating and using iterators.\n\
4680\n\
4681Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004682count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004683cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004684repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004685\n\
4686Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004687accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004688chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4689chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004690compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4691dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4692groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004693filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004694islice(seq, [start,] stop [, step]) --> elements from\n\
4695 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004696starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004697tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004698takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004699zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004700\n\
4701Combinatoric generators:\n\
4702product(p, q, ... [repeat=1]) --> cartesian product\n\
4703permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004704combinations(p, r)\n\
4705combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004706");
4707
4708
Raymond Hettingerad983e72003-11-12 14:32:26 +00004709static PyMethodDef module_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03004710 ITERTOOLS_TEE_METHODDEF
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004711 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004712};
4713
Martin v. Löwis1a214512008-06-11 05:26:20 +00004714
4715static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004716 PyModuleDef_HEAD_INIT,
4717 "itertools",
4718 module_doc,
4719 -1,
4720 module_methods,
4721 NULL,
4722 NULL,
4723 NULL,
4724 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004725};
4726
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004727PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004728PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004729{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004730 int i;
4731 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004732 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004733 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004734 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004735 &combinations_type,
4736 &cwr_type,
4737 &cycle_type,
4738 &dropwhile_type,
4739 &takewhile_type,
4740 &islice_type,
4741 &starmap_type,
4742 &chain_type,
4743 &compress_type,
4744 &filterfalse_type,
4745 &count_type,
4746 &ziplongest_type,
4747 &permutations_type,
4748 &product_type,
4749 &repeat_type,
4750 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004751 &_grouper_type,
4752 &tee_type,
4753 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004754 NULL
4755 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004757 Py_TYPE(&teedataobject_type) = &PyType_Type;
4758 m = PyModule_Create(&itertoolsmodule);
4759 if (m == NULL)
4760 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004761
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004762 for (i=0 ; typelist[i] != NULL ; i++) {
4763 if (PyType_Ready(typelist[i]) < 0)
4764 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004765 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004766 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004767 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004768 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004769
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004770 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004771}