blob: 985915f09d56830d0d9b4e9e0d0bf9b2fc235108 [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
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000010
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070011/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000012
13typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000014 PyObject_HEAD
15 PyObject *it;
16 PyObject *keyfunc;
17 PyObject *tgtkey;
18 PyObject *currkey;
19 PyObject *currvalue;
Serhiy Storchakac247caf2017-09-24 13:36:11 +030020 const void *currgrouper; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
Serhiy Storchakac740e4f2017-09-26 21:47:56 +030076Py_LOCAL_INLINE(int)
77groupby_step(groupbyobject *gbo)
78{
79 PyObject *newvalue, *newkey, *oldvalue;
80
81 newvalue = PyIter_Next(gbo->it);
82 if (newvalue == NULL)
83 return -1;
84
85 if (gbo->keyfunc == Py_None) {
86 newkey = newvalue;
87 Py_INCREF(newvalue);
88 } else {
89 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
90 if (newkey == NULL) {
91 Py_DECREF(newvalue);
92 return -1;
93 }
94 }
95
96 oldvalue = gbo->currvalue;
97 gbo->currvalue = newvalue;
98 Py_XSETREF(gbo->currkey, newkey);
99 Py_XDECREF(oldvalue);
100 return 0;
101}
102
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000103static PyObject *
104groupby_next(groupbyobject *gbo)
105{
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300106 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000107
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300108 gbo->currgrouper = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000109 /* skip to next iteration group */
110 for (;;) {
111 if (gbo->currkey == NULL)
112 /* pass */;
113 else if (gbo->tgtkey == NULL)
114 break;
115 else {
116 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000117
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700118 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000119 if (rcmp == -1)
120 return NULL;
121 else if (rcmp == 0)
122 break;
123 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000124
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300125 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000126 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000127 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300129 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000130
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000131 grouper = _grouper_create(gbo, gbo->tgtkey);
132 if (grouper == NULL)
133 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000134
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000135 r = PyTuple_Pack(2, gbo->currkey, grouper);
136 Py_DECREF(grouper);
137 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000138}
139
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000140static PyObject *
141groupby_reduce(groupbyobject *lz)
142{
143 /* reduce as a 'new' call with an optional 'setstate' if groupby
144 * has started
145 */
146 PyObject *value;
147 if (lz->tgtkey && lz->currkey && lz->currvalue)
148 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
149 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
150 else
151 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
152 lz->it, lz->keyfunc);
153
154 return value;
155}
156
157PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
158
159static PyObject *
160groupby_setstate(groupbyobject *lz, PyObject *state)
161{
162 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300163 if (!PyTuple_Check(state)) {
164 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000165 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300166 }
167 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
168 return NULL;
169 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200170 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300171 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200172 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300173 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200174 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300175 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000176 Py_RETURN_NONE;
177}
178
179PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
180
181static PyMethodDef groupby_methods[] = {
182 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
183 reduce_doc},
184 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
185 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700186 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000187};
188
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000189PyDoc_STRVAR(groupby_doc,
Raymond Hettinger49392c62017-09-25 01:21:06 -0700190"groupby(iterable, key=None) -> make an iterator that returns consecutive\n\
191keys and groups from the iterable. If the key function is not specified or\n\
192is None, the element itself is used for grouping.\n");
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000193
194static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000195 PyVarObject_HEAD_INIT(NULL, 0)
196 "itertools.groupby", /* tp_name */
197 sizeof(groupbyobject), /* tp_basicsize */
198 0, /* tp_itemsize */
199 /* methods */
200 (destructor)groupby_dealloc, /* tp_dealloc */
201 0, /* tp_print */
202 0, /* tp_getattr */
203 0, /* tp_setattr */
204 0, /* tp_reserved */
205 0, /* tp_repr */
206 0, /* tp_as_number */
207 0, /* tp_as_sequence */
208 0, /* tp_as_mapping */
209 0, /* tp_hash */
210 0, /* tp_call */
211 0, /* tp_str */
212 PyObject_GenericGetAttr, /* tp_getattro */
213 0, /* tp_setattro */
214 0, /* tp_as_buffer */
215 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
216 Py_TPFLAGS_BASETYPE, /* tp_flags */
217 groupby_doc, /* tp_doc */
218 (traverseproc)groupby_traverse, /* tp_traverse */
219 0, /* tp_clear */
220 0, /* tp_richcompare */
221 0, /* tp_weaklistoffset */
222 PyObject_SelfIter, /* tp_iter */
223 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000224 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000225 0, /* tp_members */
226 0, /* tp_getset */
227 0, /* tp_base */
228 0, /* tp_dict */
229 0, /* tp_descr_get */
230 0, /* tp_descr_set */
231 0, /* tp_dictoffset */
232 0, /* tp_init */
233 0, /* tp_alloc */
234 groupby_new, /* tp_new */
235 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236};
237
238
239/* _grouper object (internal) ************************************************/
240
241typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000242 PyObject_HEAD
243 PyObject *parent;
244 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000245} _grouperobject;
246
247static PyTypeObject _grouper_type;
248
249static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000250_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
251{
252 PyObject *parent, *tgtkey;
253
254 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
255 return NULL;
256
257 return _grouper_create((groupbyobject*) parent, tgtkey);
258}
259
260static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000261_grouper_create(groupbyobject *parent, PyObject *tgtkey)
262{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000263 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000265 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
266 if (igo == NULL)
267 return NULL;
268 igo->parent = (PyObject *)parent;
269 Py_INCREF(parent);
270 igo->tgtkey = tgtkey;
271 Py_INCREF(tgtkey);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300272 parent->currgrouper = igo; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000273
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000274 PyObject_GC_Track(igo);
275 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000276}
277
278static void
279_grouper_dealloc(_grouperobject *igo)
280{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000281 PyObject_GC_UnTrack(igo);
282 Py_DECREF(igo->parent);
283 Py_DECREF(igo->tgtkey);
284 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000285}
286
287static int
288_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
289{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000290 Py_VISIT(igo->parent);
291 Py_VISIT(igo->tgtkey);
292 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000293}
294
295static PyObject *
296_grouper_next(_grouperobject *igo)
297{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000298 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300299 PyObject *r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000301
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300302 if (gbo->currgrouper != igo)
303 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000304 if (gbo->currvalue == NULL) {
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300305 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000306 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000307 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000308
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000309 assert(gbo->currkey != NULL);
310 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
311 if (rcmp <= 0)
312 /* got any error or current group is end */
313 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000315 r = gbo->currvalue;
316 gbo->currvalue = NULL;
317 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000318
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000319 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000320}
321
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000322static PyObject *
323_grouper_reduce(_grouperobject *lz)
324{
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300325 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
326 return Py_BuildValue("N(())", _PyObject_GetBuiltin("iter"));
327 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700328 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000329}
330
331static PyMethodDef _grouper_methods[] = {
332 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
333 reduce_doc},
334 {NULL, NULL} /* sentinel */
335};
336
337
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000338static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000339 PyVarObject_HEAD_INIT(NULL, 0)
340 "itertools._grouper", /* tp_name */
341 sizeof(_grouperobject), /* tp_basicsize */
342 0, /* tp_itemsize */
343 /* methods */
344 (destructor)_grouper_dealloc, /* tp_dealloc */
345 0, /* tp_print */
346 0, /* tp_getattr */
347 0, /* tp_setattr */
348 0, /* tp_reserved */
349 0, /* tp_repr */
350 0, /* tp_as_number */
351 0, /* tp_as_sequence */
352 0, /* tp_as_mapping */
353 0, /* tp_hash */
354 0, /* tp_call */
355 0, /* tp_str */
356 PyObject_GenericGetAttr, /* tp_getattro */
357 0, /* tp_setattro */
358 0, /* tp_as_buffer */
359 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
360 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700361 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000362 0, /* tp_clear */
363 0, /* tp_richcompare */
364 0, /* tp_weaklistoffset */
365 PyObject_SelfIter, /* tp_iter */
366 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000367 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000368 0, /* tp_members */
369 0, /* tp_getset */
370 0, /* tp_base */
371 0, /* tp_dict */
372 0, /* tp_descr_get */
373 0, /* tp_descr_set */
374 0, /* tp_dictoffset */
375 0, /* tp_init */
376 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000377 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000378 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000379};
380
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000381
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700382/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000383
Raymond Hettingerad983e72003-11-12 14:32:26 +0000384/* The teedataobject pre-allocates space for LINKCELLS number of objects.
385 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000386 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000387 the other structure members including PyHEAD overhead. The larger the
388 value, the less memory overhead per object and the less time spent
389 allocating/deallocating new links. The smaller the number, the less
390 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000392#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000393
394typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000395 PyObject_HEAD
396 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700397 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000398 PyObject *nextlink;
399 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000400} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000401
402typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000403 PyObject_HEAD
404 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700405 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000406 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000407} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000408
Raymond Hettingerad983e72003-11-12 14:32:26 +0000409static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000410
411static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000412teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000413{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000414 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000416 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
417 if (tdo == NULL)
418 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000419
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000420 tdo->numread = 0;
421 tdo->nextlink = NULL;
422 Py_INCREF(it);
423 tdo->it = it;
424 PyObject_GC_Track(tdo);
425 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426}
427
428static PyObject *
429teedataobject_jumplink(teedataobject *tdo)
430{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000431 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000432 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000433 Py_XINCREF(tdo->nextlink);
434 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000435}
436
437static PyObject *
438teedataobject_getitem(teedataobject *tdo, int i)
439{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000440 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000441
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000442 assert(i < LINKCELLS);
443 if (i < tdo->numread)
444 value = tdo->values[i];
445 else {
446 /* this is the lead iterator, so fetch more data */
447 assert(i == tdo->numread);
448 value = PyIter_Next(tdo->it);
449 if (value == NULL)
450 return NULL;
451 tdo->numread++;
452 tdo->values[i] = value;
453 }
454 Py_INCREF(value);
455 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000456}
457
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000458static int
459teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000461 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 Py_VISIT(tdo->it);
464 for (i = 0; i < tdo->numread; i++)
465 Py_VISIT(tdo->values[i]);
466 Py_VISIT(tdo->nextlink);
467 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000468}
469
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200470static void
471teedataobject_safe_decref(PyObject *obj)
472{
473 while (obj && Py_TYPE(obj) == &teedataobject_type &&
474 Py_REFCNT(obj) == 1) {
475 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
476 ((teedataobject *)obj)->nextlink = NULL;
477 Py_DECREF(obj);
478 obj = nextlink;
479 }
480 Py_XDECREF(obj);
481}
482
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000483static int
484teedataobject_clear(teedataobject *tdo)
485{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000486 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200487 PyObject *tmp;
488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000489 Py_CLEAR(tdo->it);
490 for (i=0 ; i<tdo->numread ; i++)
491 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200492 tmp = tdo->nextlink;
493 tdo->nextlink = NULL;
494 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000495 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000496}
497
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000498static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000499teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000500{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 PyObject_GC_UnTrack(tdo);
502 teedataobject_clear(tdo);
503 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504}
505
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000506static PyObject *
507teedataobject_reduce(teedataobject *tdo)
508{
509 int i;
510 /* create a temporary list of already iterated values */
511 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700512
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000513 if (!values)
514 return NULL;
515 for (i=0 ; i<tdo->numread ; i++) {
516 Py_INCREF(tdo->values[i]);
517 PyList_SET_ITEM(values, i, tdo->values[i]);
518 }
519 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
520 values,
521 tdo->nextlink ? tdo->nextlink : Py_None);
522}
523
524static PyTypeObject teedataobject_type;
525
526static PyObject *
527teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
528{
529 teedataobject *tdo;
530 PyObject *it, *values, *next;
531 Py_ssize_t i, len;
532
533 assert(type == &teedataobject_type);
534 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
535 return NULL;
536
537 tdo = (teedataobject *)teedataobject_newinternal(it);
538 if (!tdo)
539 return NULL;
540
541 len = PyList_GET_SIZE(values);
542 if (len > LINKCELLS)
543 goto err;
544 for (i=0; i<len; i++) {
545 tdo->values[i] = PyList_GET_ITEM(values, i);
546 Py_INCREF(tdo->values[i]);
547 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200548 /* len <= LINKCELLS < INT_MAX */
549 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000550
551 if (len == LINKCELLS) {
552 if (next != Py_None) {
553 if (Py_TYPE(next) != &teedataobject_type)
554 goto err;
555 assert(tdo->nextlink == NULL);
556 Py_INCREF(next);
557 tdo->nextlink = next;
558 }
559 } else {
560 if (next != Py_None)
561 goto err; /* shouldn't have a next if we are not full */
562 }
563 return (PyObject*)tdo;
564
565err:
566 Py_XDECREF(tdo);
567 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
568 return NULL;
569}
570
571static PyMethodDef teedataobject_methods[] = {
572 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
573 reduce_doc},
574 {NULL, NULL} /* sentinel */
575};
576
Raymond Hettingerad983e72003-11-12 14:32:26 +0000577PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000578
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700580 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000581 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000582 sizeof(teedataobject), /* tp_basicsize */
583 0, /* tp_itemsize */
584 /* methods */
585 (destructor)teedataobject_dealloc, /* tp_dealloc */
586 0, /* tp_print */
587 0, /* tp_getattr */
588 0, /* tp_setattr */
589 0, /* tp_reserved */
590 0, /* tp_repr */
591 0, /* tp_as_number */
592 0, /* tp_as_sequence */
593 0, /* tp_as_mapping */
594 0, /* tp_hash */
595 0, /* tp_call */
596 0, /* tp_str */
597 PyObject_GenericGetAttr, /* tp_getattro */
598 0, /* tp_setattro */
599 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700600 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000601 teedataobject_doc, /* tp_doc */
602 (traverseproc)teedataobject_traverse, /* tp_traverse */
603 (inquiry)teedataobject_clear, /* tp_clear */
604 0, /* tp_richcompare */
605 0, /* tp_weaklistoffset */
606 0, /* tp_iter */
607 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000608 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000609 0, /* tp_members */
610 0, /* tp_getset */
611 0, /* tp_base */
612 0, /* tp_dict */
613 0, /* tp_descr_get */
614 0, /* tp_descr_set */
615 0, /* tp_dictoffset */
616 0, /* tp_init */
617 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000618 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000619 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000620};
621
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622
623static PyTypeObject tee_type;
624
625static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000626tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000628 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000630 if (to->index >= LINKCELLS) {
631 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200632 if (link == NULL)
633 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300634 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 to->index = 0;
636 }
637 value = teedataobject_getitem(to->dataobj, to->index);
638 if (value == NULL)
639 return NULL;
640 to->index++;
641 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000642}
643
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000644static int
645tee_traverse(teeobject *to, visitproc visit, void *arg)
646{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000647 Py_VISIT((PyObject *)to->dataobj);
648 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000649}
650
Raymond Hettingerad983e72003-11-12 14:32:26 +0000651static PyObject *
652tee_copy(teeobject *to)
653{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000654 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000655
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000656 newto = PyObject_GC_New(teeobject, &tee_type);
657 if (newto == NULL)
658 return NULL;
659 Py_INCREF(to->dataobj);
660 newto->dataobj = to->dataobj;
661 newto->index = to->index;
662 newto->weakreflist = NULL;
663 PyObject_GC_Track(newto);
664 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
666
667PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
668
669static PyObject *
670tee_fromiterable(PyObject *iterable)
671{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000672 teeobject *to;
673 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000675 it = PyObject_GetIter(iterable);
676 if (it == NULL)
677 return NULL;
678 if (PyObject_TypeCheck(it, &tee_type)) {
679 to = (teeobject *)tee_copy((teeobject *)it);
680 goto done;
681 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000682
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000683 to = PyObject_GC_New(teeobject, &tee_type);
684 if (to == NULL)
685 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000686 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 if (!to->dataobj) {
688 PyObject_GC_Del(to);
689 to = NULL;
690 goto done;
691 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 to->index = 0;
694 to->weakreflist = NULL;
695 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000696done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000697 Py_XDECREF(it);
698 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000699}
700
701static PyObject *
702tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
703{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000704 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000705
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000706 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000707 return NULL;
708 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000709}
710
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000711static int
712tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000714 if (to->weakreflist != NULL)
715 PyObject_ClearWeakRefs((PyObject *) to);
716 Py_CLEAR(to->dataobj);
717 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000718}
719
720static void
721tee_dealloc(teeobject *to)
722{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000723 PyObject_GC_UnTrack(to);
724 tee_clear(to);
725 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000726}
727
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000728static PyObject *
729tee_reduce(teeobject *to)
730{
731 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
732}
733
734static PyObject *
735tee_setstate(teeobject *to, PyObject *state)
736{
737 teedataobject *tdo;
738 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300739 if (!PyTuple_Check(state)) {
740 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000741 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300742 }
743 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
744 return NULL;
745 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000746 if (index < 0 || index > LINKCELLS) {
747 PyErr_SetString(PyExc_ValueError, "Index out of range");
748 return NULL;
749 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200750 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300751 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000752 to->index = index;
753 Py_RETURN_NONE;
754}
755
Raymond Hettingerad983e72003-11-12 14:32:26 +0000756PyDoc_STRVAR(teeobject_doc,
757"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000758
Raymond Hettingerad983e72003-11-12 14:32:26 +0000759static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700760 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
761 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
762 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000763 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000764};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000765
766static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000767 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000768 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 sizeof(teeobject), /* tp_basicsize */
770 0, /* tp_itemsize */
771 /* methods */
772 (destructor)tee_dealloc, /* tp_dealloc */
773 0, /* tp_print */
774 0, /* tp_getattr */
775 0, /* tp_setattr */
776 0, /* tp_reserved */
777 0, /* tp_repr */
778 0, /* tp_as_number */
779 0, /* tp_as_sequence */
780 0, /* tp_as_mapping */
781 0, /* tp_hash */
782 0, /* tp_call */
783 0, /* tp_str */
784 0, /* tp_getattro */
785 0, /* tp_setattro */
786 0, /* tp_as_buffer */
787 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
788 teeobject_doc, /* tp_doc */
789 (traverseproc)tee_traverse, /* tp_traverse */
790 (inquiry)tee_clear, /* tp_clear */
791 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700792 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000793 PyObject_SelfIter, /* tp_iter */
794 (iternextfunc)tee_next, /* tp_iternext */
795 tee_methods, /* tp_methods */
796 0, /* tp_members */
797 0, /* tp_getset */
798 0, /* tp_base */
799 0, /* tp_dict */
800 0, /* tp_descr_get */
801 0, /* tp_descr_set */
802 0, /* tp_dictoffset */
803 0, /* tp_init */
804 0, /* tp_alloc */
805 tee_new, /* tp_new */
806 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000807};
808
Raymond Hettingerad983e72003-11-12 14:32:26 +0000809static PyObject *
810tee(PyObject *self, PyObject *args)
811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000812 Py_ssize_t i, n=2;
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200813 PyObject *it, *iterable, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200814 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000816 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
817 return NULL;
818 if (n < 0) {
819 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
820 return NULL;
821 }
822 result = PyTuple_New(n);
823 if (result == NULL)
824 return NULL;
825 if (n == 0)
826 return result;
827 it = PyObject_GetIter(iterable);
828 if (it == NULL) {
829 Py_DECREF(result);
830 return NULL;
831 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200832
833 copyfunc = _PyObject_GetAttrId(it, &PyId___copy__);
834 if (copyfunc != NULL) {
835 copyable = it;
836 }
837 else if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
838 Py_DECREF(it);
839 Py_DECREF(result);
840 return NULL;
841 }
842 else {
843 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000844 copyable = tee_fromiterable(it);
845 Py_DECREF(it);
846 if (copyable == NULL) {
847 Py_DECREF(result);
848 return NULL;
849 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200850 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
851 if (copyfunc == NULL) {
852 Py_DECREF(copyable);
853 Py_DECREF(result);
854 return NULL;
855 }
856 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200857
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200858 PyTuple_SET_ITEM(result, 0, copyable);
859 for (i = 1; i < n; i++) {
860 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000861 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200862 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 Py_DECREF(result);
864 return NULL;
865 }
866 PyTuple_SET_ITEM(result, i, copyable);
867 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200868 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000869 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000870}
871
872PyDoc_STRVAR(tee_doc,
873"tee(iterable, n=2) --> tuple of n independent iterators.");
874
875
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700876/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000877
878typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000879 PyObject_HEAD
880 PyObject *it;
881 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700882 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000883 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000884} cycleobject;
885
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000886static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000887
888static PyObject *
889cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
890{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000891 PyObject *it;
892 PyObject *iterable;
893 PyObject *saved;
894 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000895
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +0300896 if (type == &cycle_type && !_PyArg_NoKeywords("cycle", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000897 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000899 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
900 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000901
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000902 /* Get iterator. */
903 it = PyObject_GetIter(iterable);
904 if (it == NULL)
905 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000907 saved = PyList_New(0);
908 if (saved == NULL) {
909 Py_DECREF(it);
910 return NULL;
911 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 /* create cycleobject structure */
914 lz = (cycleobject *)type->tp_alloc(type, 0);
915 if (lz == NULL) {
916 Py_DECREF(it);
917 Py_DECREF(saved);
918 return NULL;
919 }
920 lz->it = it;
921 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700922 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000923 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000925 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000926}
927
928static void
929cycle_dealloc(cycleobject *lz)
930{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000931 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000932 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700933 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000935}
936
937static int
938cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
939{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700940 if (lz->it)
941 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000942 Py_VISIT(lz->saved);
943 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000944}
945
946static PyObject *
947cycle_next(cycleobject *lz)
948{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000949 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000950
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700951 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000952 item = PyIter_Next(lz->it);
953 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700954 if (lz->firstpass)
955 return item;
956 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000957 Py_DECREF(item);
958 return NULL;
959 }
960 return item;
961 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -0700962 /* Note: StopIteration is already cleared by PyIter_Next() */
963 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000964 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700965 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000966 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +0200967 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700968 return NULL;
969 item = PyList_GET_ITEM(lz->saved, lz->index);
970 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +0200971 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700972 lz->index = 0;
973 Py_INCREF(item);
974 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000975}
976
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000977static PyObject *
978cycle_reduce(cycleobject *lz)
979{
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700980 /* Create a new cycle with the iterator tuple, then set the saved state */
981 if (lz->it == NULL) {
982 PyObject *it = PyObject_GetIter(lz->saved);
983 if (it == NULL)
984 return NULL;
985 if (lz->index != 0) {
986 _Py_IDENTIFIER(__setstate__);
987 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
988 "n", lz->index);
989 if (res == NULL) {
990 Py_DECREF(it);
991 return NULL;
992 }
993 Py_DECREF(res);
994 }
995 return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
996 }
997 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
998 lz->firstpass);
Raymond Hettingerc39786d2015-08-15 15:16:12 -0700999}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001000
1001static PyObject *
1002cycle_setstate(cycleobject *lz, PyObject *state)
1003{
1004 PyObject *saved=NULL;
1005 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001006 if (!PyTuple_Check(state)) {
1007 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001008 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001009 }
1010 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1011 return NULL;
1012 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001013 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001014 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001015 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001016 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001017 Py_RETURN_NONE;
1018}
1019
1020static PyMethodDef cycle_methods[] = {
1021 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1022 reduce_doc},
1023 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1024 setstate_doc},
1025 {NULL, NULL} /* sentinel */
1026};
1027
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001028PyDoc_STRVAR(cycle_doc,
1029"cycle(iterable) --> cycle object\n\
1030\n\
1031Return elements from the iterable until it is exhausted.\n\
1032Then repeat the sequence indefinitely.");
1033
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001034static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001035 PyVarObject_HEAD_INIT(NULL, 0)
1036 "itertools.cycle", /* tp_name */
1037 sizeof(cycleobject), /* tp_basicsize */
1038 0, /* tp_itemsize */
1039 /* methods */
1040 (destructor)cycle_dealloc, /* tp_dealloc */
1041 0, /* tp_print */
1042 0, /* tp_getattr */
1043 0, /* tp_setattr */
1044 0, /* tp_reserved */
1045 0, /* tp_repr */
1046 0, /* tp_as_number */
1047 0, /* tp_as_sequence */
1048 0, /* tp_as_mapping */
1049 0, /* tp_hash */
1050 0, /* tp_call */
1051 0, /* tp_str */
1052 PyObject_GenericGetAttr, /* tp_getattro */
1053 0, /* tp_setattro */
1054 0, /* tp_as_buffer */
1055 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1056 Py_TPFLAGS_BASETYPE, /* tp_flags */
1057 cycle_doc, /* tp_doc */
1058 (traverseproc)cycle_traverse, /* tp_traverse */
1059 0, /* tp_clear */
1060 0, /* tp_richcompare */
1061 0, /* tp_weaklistoffset */
1062 PyObject_SelfIter, /* tp_iter */
1063 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001064 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 0, /* tp_members */
1066 0, /* tp_getset */
1067 0, /* tp_base */
1068 0, /* tp_dict */
1069 0, /* tp_descr_get */
1070 0, /* tp_descr_set */
1071 0, /* tp_dictoffset */
1072 0, /* tp_init */
1073 0, /* tp_alloc */
1074 cycle_new, /* tp_new */
1075 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001076};
1077
1078
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001079/* dropwhile object **********************************************************/
1080
1081typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001082 PyObject_HEAD
1083 PyObject *func;
1084 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001085 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001086} dropwhileobject;
1087
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001088static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089
1090static PyObject *
1091dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1092{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001093 PyObject *func, *seq;
1094 PyObject *it;
1095 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001096
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001097 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001098 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001100 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1101 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001102
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001103 /* Get iterator. */
1104 it = PyObject_GetIter(seq);
1105 if (it == NULL)
1106 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001107
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001108 /* create dropwhileobject structure */
1109 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1110 if (lz == NULL) {
1111 Py_DECREF(it);
1112 return NULL;
1113 }
1114 Py_INCREF(func);
1115 lz->func = func;
1116 lz->it = it;
1117 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001118
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001119 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001120}
1121
1122static void
1123dropwhile_dealloc(dropwhileobject *lz)
1124{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001125 PyObject_GC_UnTrack(lz);
1126 Py_XDECREF(lz->func);
1127 Py_XDECREF(lz->it);
1128 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001129}
1130
1131static int
1132dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1133{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 Py_VISIT(lz->it);
1135 Py_VISIT(lz->func);
1136 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001137}
1138
1139static PyObject *
1140dropwhile_next(dropwhileobject *lz)
1141{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001142 PyObject *item, *good;
1143 PyObject *it = lz->it;
1144 long ok;
1145 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001146
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001147 iternext = *Py_TYPE(it)->tp_iternext;
1148 for (;;) {
1149 item = iternext(it);
1150 if (item == NULL)
1151 return NULL;
1152 if (lz->start == 1)
1153 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001155 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001156 if (good == NULL) {
1157 Py_DECREF(item);
1158 return NULL;
1159 }
1160 ok = PyObject_IsTrue(good);
1161 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001162 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001163 lz->start = 1;
1164 return item;
1165 }
1166 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001167 if (ok < 0)
1168 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001169 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170}
1171
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001172static PyObject *
1173dropwhile_reduce(dropwhileobject *lz)
1174{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001175 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 +00001176}
1177
1178static PyObject *
1179dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1180{
1181 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001182 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001183 return NULL;
1184 lz->start = start;
1185 Py_RETURN_NONE;
1186}
1187
1188static PyMethodDef dropwhile_methods[] = {
1189 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1190 reduce_doc},
1191 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1192 setstate_doc},
1193 {NULL, NULL} /* sentinel */
1194};
1195
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196PyDoc_STRVAR(dropwhile_doc,
1197"dropwhile(predicate, iterable) --> dropwhile object\n\
1198\n\
1199Drop items from the iterable while predicate(item) is true.\n\
1200Afterwards, return every element until the iterable is exhausted.");
1201
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001202static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001203 PyVarObject_HEAD_INIT(NULL, 0)
1204 "itertools.dropwhile", /* tp_name */
1205 sizeof(dropwhileobject), /* tp_basicsize */
1206 0, /* tp_itemsize */
1207 /* methods */
1208 (destructor)dropwhile_dealloc, /* tp_dealloc */
1209 0, /* tp_print */
1210 0, /* tp_getattr */
1211 0, /* tp_setattr */
1212 0, /* tp_reserved */
1213 0, /* tp_repr */
1214 0, /* tp_as_number */
1215 0, /* tp_as_sequence */
1216 0, /* tp_as_mapping */
1217 0, /* tp_hash */
1218 0, /* tp_call */
1219 0, /* tp_str */
1220 PyObject_GenericGetAttr, /* tp_getattro */
1221 0, /* tp_setattro */
1222 0, /* tp_as_buffer */
1223 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1224 Py_TPFLAGS_BASETYPE, /* tp_flags */
1225 dropwhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001226 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001227 0, /* tp_clear */
1228 0, /* tp_richcompare */
1229 0, /* tp_weaklistoffset */
1230 PyObject_SelfIter, /* tp_iter */
1231 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001232 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001233 0, /* tp_members */
1234 0, /* tp_getset */
1235 0, /* tp_base */
1236 0, /* tp_dict */
1237 0, /* tp_descr_get */
1238 0, /* tp_descr_set */
1239 0, /* tp_dictoffset */
1240 0, /* tp_init */
1241 0, /* tp_alloc */
1242 dropwhile_new, /* tp_new */
1243 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244};
1245
1246
1247/* takewhile object **********************************************************/
1248
1249typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001250 PyObject_HEAD
1251 PyObject *func;
1252 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001253 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254} takewhileobject;
1255
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001256static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001257
1258static PyObject *
1259takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1260{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001261 PyObject *func, *seq;
1262 PyObject *it;
1263 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001264
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001265 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001267
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001268 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1269 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001270
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001271 /* Get iterator. */
1272 it = PyObject_GetIter(seq);
1273 if (it == NULL)
1274 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001276 /* create takewhileobject structure */
1277 lz = (takewhileobject *)type->tp_alloc(type, 0);
1278 if (lz == NULL) {
1279 Py_DECREF(it);
1280 return NULL;
1281 }
1282 Py_INCREF(func);
1283 lz->func = func;
1284 lz->it = it;
1285 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001288}
1289
1290static void
1291takewhile_dealloc(takewhileobject *lz)
1292{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001293 PyObject_GC_UnTrack(lz);
1294 Py_XDECREF(lz->func);
1295 Py_XDECREF(lz->it);
1296 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001297}
1298
1299static int
1300takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1301{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001302 Py_VISIT(lz->it);
1303 Py_VISIT(lz->func);
1304 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001305}
1306
1307static PyObject *
1308takewhile_next(takewhileobject *lz)
1309{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001310 PyObject *item, *good;
1311 PyObject *it = lz->it;
1312 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001313
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001314 if (lz->stop == 1)
1315 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001317 item = (*Py_TYPE(it)->tp_iternext)(it);
1318 if (item == NULL)
1319 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001320
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01001321 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001322 if (good == NULL) {
1323 Py_DECREF(item);
1324 return NULL;
1325 }
1326 ok = PyObject_IsTrue(good);
1327 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001328 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001329 return item;
1330 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001331 if (ok == 0)
1332 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001333 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334}
1335
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001336static PyObject *
1337takewhile_reduce(takewhileobject *lz)
1338{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001339 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 +00001340}
1341
1342static PyObject *
1343takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1344{
1345 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001346
Antoine Pitrou721738f2012-08-15 23:20:39 +02001347 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001348 return NULL;
1349 lz->stop = stop;
1350 Py_RETURN_NONE;
1351}
1352
1353static PyMethodDef takewhile_reduce_methods[] = {
1354 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1355 reduce_doc},
1356 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1357 setstate_doc},
1358 {NULL, NULL} /* sentinel */
1359};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360PyDoc_STRVAR(takewhile_doc,
1361"takewhile(predicate, iterable) --> takewhile object\n\
1362\n\
1363Return successive entries from an iterable as long as the \n\
1364predicate evaluates to true for each entry.");
1365
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001366static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001367 PyVarObject_HEAD_INIT(NULL, 0)
1368 "itertools.takewhile", /* tp_name */
1369 sizeof(takewhileobject), /* tp_basicsize */
1370 0, /* tp_itemsize */
1371 /* methods */
1372 (destructor)takewhile_dealloc, /* tp_dealloc */
1373 0, /* tp_print */
1374 0, /* tp_getattr */
1375 0, /* tp_setattr */
1376 0, /* tp_reserved */
1377 0, /* tp_repr */
1378 0, /* tp_as_number */
1379 0, /* tp_as_sequence */
1380 0, /* tp_as_mapping */
1381 0, /* tp_hash */
1382 0, /* tp_call */
1383 0, /* tp_str */
1384 PyObject_GenericGetAttr, /* tp_getattro */
1385 0, /* tp_setattro */
1386 0, /* tp_as_buffer */
1387 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1388 Py_TPFLAGS_BASETYPE, /* tp_flags */
1389 takewhile_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001390 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001391 0, /* tp_clear */
1392 0, /* tp_richcompare */
1393 0, /* tp_weaklistoffset */
1394 PyObject_SelfIter, /* tp_iter */
1395 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001396 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001397 0, /* tp_members */
1398 0, /* tp_getset */
1399 0, /* tp_base */
1400 0, /* tp_dict */
1401 0, /* tp_descr_get */
1402 0, /* tp_descr_set */
1403 0, /* tp_dictoffset */
1404 0, /* tp_init */
1405 0, /* tp_alloc */
1406 takewhile_new, /* tp_new */
1407 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001408};
1409
1410
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001411/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001412
1413typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001414 PyObject_HEAD
1415 PyObject *it;
1416 Py_ssize_t next;
1417 Py_ssize_t stop;
1418 Py_ssize_t step;
1419 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001420} isliceobject;
1421
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001422static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423
1424static PyObject *
1425islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1426{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001427 PyObject *seq;
1428 Py_ssize_t start=0, stop=-1, step=1;
1429 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1430 Py_ssize_t numargs;
1431 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001433 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001434 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001435
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001436 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1437 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001439 numargs = PyTuple_Size(args);
1440 if (numargs == 2) {
1441 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001442 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001443 if (stop == -1) {
1444 if (PyErr_Occurred())
1445 PyErr_Clear();
1446 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001447 "Stop argument for islice() must be None or "
1448 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001449 return NULL;
1450 }
1451 }
1452 } else {
1453 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001454 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001455 if (start == -1 && PyErr_Occurred())
1456 PyErr_Clear();
1457 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001458 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001459 if (stop == -1) {
1460 if (PyErr_Occurred())
1461 PyErr_Clear();
1462 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001463 "Stop argument for islice() must be None or "
1464 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001465 return NULL;
1466 }
1467 }
1468 }
1469 if (start<0 || stop<-1) {
1470 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001471 "Indices for islice() must be None or "
1472 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001473 return NULL;
1474 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001476 if (a3 != NULL) {
1477 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001478 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001479 if (step == -1 && PyErr_Occurred())
1480 PyErr_Clear();
1481 }
1482 if (step<1) {
1483 PyErr_SetString(PyExc_ValueError,
1484 "Step for islice() must be a positive integer or None.");
1485 return NULL;
1486 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001488 /* Get iterator. */
1489 it = PyObject_GetIter(seq);
1490 if (it == NULL)
1491 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001492
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001493 /* create isliceobject structure */
1494 lz = (isliceobject *)type->tp_alloc(type, 0);
1495 if (lz == NULL) {
1496 Py_DECREF(it);
1497 return NULL;
1498 }
1499 lz->it = it;
1500 lz->next = start;
1501 lz->stop = stop;
1502 lz->step = step;
1503 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001505 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001506}
1507
1508static void
1509islice_dealloc(isliceobject *lz)
1510{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001511 PyObject_GC_UnTrack(lz);
1512 Py_XDECREF(lz->it);
1513 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001514}
1515
1516static int
1517islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1518{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001519 Py_VISIT(lz->it);
1520 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001521}
1522
1523static PyObject *
1524islice_next(isliceobject *lz)
1525{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001526 PyObject *item;
1527 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001528 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001529 Py_ssize_t oldnext;
1530 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001531
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001532 if (it == NULL)
1533 return NULL;
1534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001535 iternext = *Py_TYPE(it)->tp_iternext;
1536 while (lz->cnt < lz->next) {
1537 item = iternext(it);
1538 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001539 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001540 Py_DECREF(item);
1541 lz->cnt++;
1542 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001543 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001544 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001545 item = iternext(it);
1546 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001547 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001548 lz->cnt++;
1549 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001550 /* The (size_t) cast below avoids the danger of undefined
1551 behaviour from signed integer overflow. */
1552 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001553 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1554 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001555 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001556
1557empty:
1558 Py_CLEAR(lz->it);
1559 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001560}
1561
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001562static PyObject *
1563islice_reduce(isliceobject *lz)
1564{
1565 /* When unpickled, generate a new object with the same bounds,
1566 * then 'setstate' with the next and count
1567 */
1568 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001569
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001570 if (lz->it == NULL) {
1571 PyObject *empty_list;
1572 PyObject *empty_it;
1573 empty_list = PyList_New(0);
1574 if (empty_list == NULL)
1575 return NULL;
1576 empty_it = PyObject_GetIter(empty_list);
1577 Py_DECREF(empty_list);
1578 if (empty_it == NULL)
1579 return NULL;
1580 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1581 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001582 if (lz->stop == -1) {
1583 stop = Py_None;
1584 Py_INCREF(stop);
1585 } else {
1586 stop = PyLong_FromSsize_t(lz->stop);
1587 if (stop == NULL)
1588 return NULL;
1589 }
1590 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1591 lz->it, lz->next, stop, lz->step,
1592 lz->cnt);
1593}
1594
1595static PyObject *
1596islice_setstate(isliceobject *lz, PyObject *state)
1597{
1598 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001599
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001600 if (cnt == -1 && PyErr_Occurred())
1601 return NULL;
1602 lz->cnt = cnt;
1603 Py_RETURN_NONE;
1604}
1605
1606static PyMethodDef islice_methods[] = {
1607 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1608 reduce_doc},
1609 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1610 setstate_doc},
1611 {NULL, NULL} /* sentinel */
1612};
1613
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001614PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001615"islice(iterable, stop) --> islice object\n\
1616islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617\n\
1618Return an iterator whose next() method returns selected values from an\n\
1619iterable. If start is specified, will skip all preceding elements;\n\
1620otherwise, start defaults to zero. Step defaults to one. If\n\
1621specified as another value, step determines how many values are \n\
1622skipped between successive calls. Works like a slice() on a list\n\
1623but returns an iterator.");
1624
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001625static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001626 PyVarObject_HEAD_INIT(NULL, 0)
1627 "itertools.islice", /* tp_name */
1628 sizeof(isliceobject), /* tp_basicsize */
1629 0, /* tp_itemsize */
1630 /* methods */
1631 (destructor)islice_dealloc, /* tp_dealloc */
1632 0, /* tp_print */
1633 0, /* tp_getattr */
1634 0, /* tp_setattr */
1635 0, /* tp_reserved */
1636 0, /* tp_repr */
1637 0, /* tp_as_number */
1638 0, /* tp_as_sequence */
1639 0, /* tp_as_mapping */
1640 0, /* tp_hash */
1641 0, /* tp_call */
1642 0, /* tp_str */
1643 PyObject_GenericGetAttr, /* tp_getattro */
1644 0, /* tp_setattro */
1645 0, /* tp_as_buffer */
1646 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1647 Py_TPFLAGS_BASETYPE, /* tp_flags */
1648 islice_doc, /* tp_doc */
1649 (traverseproc)islice_traverse, /* tp_traverse */
1650 0, /* tp_clear */
1651 0, /* tp_richcompare */
1652 0, /* tp_weaklistoffset */
1653 PyObject_SelfIter, /* tp_iter */
1654 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001655 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001656 0, /* tp_members */
1657 0, /* tp_getset */
1658 0, /* tp_base */
1659 0, /* tp_dict */
1660 0, /* tp_descr_get */
1661 0, /* tp_descr_set */
1662 0, /* tp_dictoffset */
1663 0, /* tp_init */
1664 0, /* tp_alloc */
1665 islice_new, /* tp_new */
1666 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001667};
1668
1669
1670/* starmap object ************************************************************/
1671
1672typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001673 PyObject_HEAD
1674 PyObject *func;
1675 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676} starmapobject;
1677
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001678static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679
1680static PyObject *
1681starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1682{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001683 PyObject *func, *seq;
1684 PyObject *it;
1685 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001686
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001687 if (type == &starmap_type && !_PyArg_NoKeywords("starmap", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001688 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001690 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1691 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001693 /* Get iterator. */
1694 it = PyObject_GetIter(seq);
1695 if (it == NULL)
1696 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001698 /* create starmapobject structure */
1699 lz = (starmapobject *)type->tp_alloc(type, 0);
1700 if (lz == NULL) {
1701 Py_DECREF(it);
1702 return NULL;
1703 }
1704 Py_INCREF(func);
1705 lz->func = func;
1706 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001707
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001708 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709}
1710
1711static void
1712starmap_dealloc(starmapobject *lz)
1713{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001714 PyObject_GC_UnTrack(lz);
1715 Py_XDECREF(lz->func);
1716 Py_XDECREF(lz->it);
1717 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001718}
1719
1720static int
1721starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1722{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001723 Py_VISIT(lz->it);
1724 Py_VISIT(lz->func);
1725 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001726}
1727
1728static PyObject *
1729starmap_next(starmapobject *lz)
1730{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001731 PyObject *args;
1732 PyObject *result;
1733 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001735 args = (*Py_TYPE(it)->tp_iternext)(it);
1736 if (args == NULL)
1737 return NULL;
1738 if (!PyTuple_CheckExact(args)) {
1739 PyObject *newargs = PySequence_Tuple(args);
1740 Py_DECREF(args);
1741 if (newargs == NULL)
1742 return NULL;
1743 args = newargs;
1744 }
1745 result = PyObject_Call(lz->func, args, NULL);
1746 Py_DECREF(args);
1747 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001748}
1749
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001750static PyObject *
1751starmap_reduce(starmapobject *lz)
1752{
1753 /* Just pickle the iterator */
1754 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1755}
1756
1757static PyMethodDef starmap_methods[] = {
1758 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1759 reduce_doc},
1760 {NULL, NULL} /* sentinel */
1761};
1762
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001763PyDoc_STRVAR(starmap_doc,
1764"starmap(function, sequence) --> starmap object\n\
1765\n\
1766Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001767with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001768
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001769static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001770 PyVarObject_HEAD_INIT(NULL, 0)
1771 "itertools.starmap", /* tp_name */
1772 sizeof(starmapobject), /* tp_basicsize */
1773 0, /* tp_itemsize */
1774 /* methods */
1775 (destructor)starmap_dealloc, /* tp_dealloc */
1776 0, /* tp_print */
1777 0, /* tp_getattr */
1778 0, /* tp_setattr */
1779 0, /* tp_reserved */
1780 0, /* tp_repr */
1781 0, /* tp_as_number */
1782 0, /* tp_as_sequence */
1783 0, /* tp_as_mapping */
1784 0, /* tp_hash */
1785 0, /* tp_call */
1786 0, /* tp_str */
1787 PyObject_GenericGetAttr, /* tp_getattro */
1788 0, /* tp_setattro */
1789 0, /* tp_as_buffer */
1790 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1791 Py_TPFLAGS_BASETYPE, /* tp_flags */
1792 starmap_doc, /* tp_doc */
1793 (traverseproc)starmap_traverse, /* tp_traverse */
1794 0, /* tp_clear */
1795 0, /* tp_richcompare */
1796 0, /* tp_weaklistoffset */
1797 PyObject_SelfIter, /* tp_iter */
1798 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001799 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001800 0, /* tp_members */
1801 0, /* tp_getset */
1802 0, /* tp_base */
1803 0, /* tp_dict */
1804 0, /* tp_descr_get */
1805 0, /* tp_descr_set */
1806 0, /* tp_dictoffset */
1807 0, /* tp_init */
1808 0, /* tp_alloc */
1809 starmap_new, /* tp_new */
1810 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001811};
1812
1813
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001814/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001815
1816typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001817 PyObject_HEAD
1818 PyObject *source; /* Iterator over input iterables */
1819 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001820} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001821
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001822static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001823
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001824static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001825chain_new_internal(PyTypeObject *type, PyObject *source)
1826{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001827 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001828
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001829 lz = (chainobject *)type->tp_alloc(type, 0);
1830 if (lz == NULL) {
1831 Py_DECREF(source);
1832 return NULL;
1833 }
1834
1835 lz->source = source;
1836 lz->active = NULL;
1837 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001838}
1839
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001840static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001841chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001842{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001843 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001844
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001845 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001846 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001848 source = PyObject_GetIter(args);
1849 if (source == NULL)
1850 return NULL;
1851
1852 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001853}
1854
1855static PyObject *
1856chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1857{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001858 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001859
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001860 source = PyObject_GetIter(arg);
1861 if (source == NULL)
1862 return NULL;
1863
1864 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001865}
1866
1867static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001868chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001869{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001870 PyObject_GC_UnTrack(lz);
1871 Py_XDECREF(lz->active);
1872 Py_XDECREF(lz->source);
1873 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001874}
1875
Raymond Hettinger2012f172003-02-07 05:32:58 +00001876static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001877chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001878{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001879 Py_VISIT(lz->source);
1880 Py_VISIT(lz->active);
1881 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001882}
1883
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001884static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001885chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001886{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001887 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001888
T. Wouters5466d4a2017-03-30 09:58:35 -07001889 /* lz->source is the iterator of iterables. If it's NULL, we've already
1890 * consumed them all. lz->active is the current iterator. If it's NULL,
1891 * we should grab a new one from lz->source. */
1892 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001893 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001894 PyObject *iterable = PyIter_Next(lz->source);
1895 if (iterable == NULL) {
1896 Py_CLEAR(lz->source);
1897 return NULL; /* no more input sources */
1898 }
1899 lz->active = PyObject_GetIter(iterable);
1900 Py_DECREF(iterable);
1901 if (lz->active == NULL) {
1902 Py_CLEAR(lz->source);
1903 return NULL; /* input not iterable */
1904 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001905 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001906 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1907 if (item != NULL)
1908 return item;
1909 if (PyErr_Occurred()) {
1910 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1911 PyErr_Clear();
1912 else
1913 return NULL; /* input raised an exception */
1914 }
1915 /* lz->active is consumed, try with the next iterable. */
1916 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001917 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001918 /* Everything had been consumed already. */
1919 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001920}
1921
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001922static PyObject *
1923chain_reduce(chainobject *lz)
1924{
1925 if (lz->source) {
1926 /* we can't pickle function objects (itertools.from_iterable) so
1927 * we must use setstate to replace the iterable. One day we
1928 * will fix pickling of functions
1929 */
1930 if (lz->active) {
1931 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1932 } else {
1933 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1934 }
1935 } else {
1936 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1937 }
1938 return NULL;
1939}
1940
1941static PyObject *
1942chain_setstate(chainobject *lz, PyObject *state)
1943{
1944 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001945
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001946 if (!PyTuple_Check(state)) {
1947 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001948 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001949 }
1950 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
1951 return NULL;
1952 }
1953 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
1954 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
1955 return NULL;
1956 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001957
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001958 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001959 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001960 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001961 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001962 Py_RETURN_NONE;
1963}
1964
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001965PyDoc_STRVAR(chain_doc,
1966"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001967\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001968Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001969first iterable until it is exhausted, then elements from the next\n\
1970iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001971
Christian Heimesf16baeb2008-02-29 14:57:44 +00001972PyDoc_STRVAR(chain_from_iterable_doc,
1973"chain.from_iterable(iterable) --> chain object\n\
1974\n\
Martin Pantereb995702016-07-28 01:11:04 +00001975Alternate chain() constructor taking a single iterable argument\n\
Christian Heimesf16baeb2008-02-29 14:57:44 +00001976that evaluates lazily.");
1977
1978static PyMethodDef chain_methods[] = {
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001979 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1980 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001981 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1982 reduce_doc},
1983 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1984 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001985 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001986};
1987
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001988static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001989 PyVarObject_HEAD_INIT(NULL, 0)
1990 "itertools.chain", /* tp_name */
1991 sizeof(chainobject), /* tp_basicsize */
1992 0, /* tp_itemsize */
1993 /* methods */
1994 (destructor)chain_dealloc, /* tp_dealloc */
1995 0, /* tp_print */
1996 0, /* tp_getattr */
1997 0, /* tp_setattr */
1998 0, /* tp_reserved */
1999 0, /* tp_repr */
2000 0, /* tp_as_number */
2001 0, /* tp_as_sequence */
2002 0, /* tp_as_mapping */
2003 0, /* tp_hash */
2004 0, /* tp_call */
2005 0, /* tp_str */
2006 PyObject_GenericGetAttr, /* tp_getattro */
2007 0, /* tp_setattro */
2008 0, /* tp_as_buffer */
2009 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2010 Py_TPFLAGS_BASETYPE, /* tp_flags */
2011 chain_doc, /* tp_doc */
2012 (traverseproc)chain_traverse, /* tp_traverse */
2013 0, /* tp_clear */
2014 0, /* tp_richcompare */
2015 0, /* tp_weaklistoffset */
2016 PyObject_SelfIter, /* tp_iter */
2017 (iternextfunc)chain_next, /* tp_iternext */
2018 chain_methods, /* tp_methods */
2019 0, /* tp_members */
2020 0, /* tp_getset */
2021 0, /* tp_base */
2022 0, /* tp_dict */
2023 0, /* tp_descr_get */
2024 0, /* tp_descr_set */
2025 0, /* tp_dictoffset */
2026 0, /* tp_init */
2027 0, /* tp_alloc */
2028 chain_new, /* tp_new */
2029 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002030};
2031
2032
Christian Heimesc3f30c42008-02-22 16:37:40 +00002033/* product object ************************************************************/
2034
2035typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002036 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002037 PyObject *pools; /* tuple of pool tuples */
2038 Py_ssize_t *indices; /* one index per pool */
2039 PyObject *result; /* most recently returned result tuple */
2040 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002041} productobject;
2042
2043static PyTypeObject product_type;
2044
2045static PyObject *
2046product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2047{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002048 productobject *lz;
2049 Py_ssize_t nargs, npools, repeat=1;
2050 PyObject *pools = NULL;
2051 Py_ssize_t *indices = NULL;
2052 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002054 if (kwds != NULL) {
2055 char *kwlist[] = {"repeat", 0};
2056 PyObject *tmpargs = PyTuple_New(0);
2057 if (tmpargs == NULL)
2058 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002059 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2060 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002061 Py_DECREF(tmpargs);
2062 return NULL;
2063 }
2064 Py_DECREF(tmpargs);
2065 if (repeat < 0) {
2066 PyErr_SetString(PyExc_ValueError,
2067 "repeat argument cannot be negative");
2068 return NULL;
2069 }
2070 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002071
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002072 assert(PyTuple_CheckExact(args));
2073 if (repeat == 0) {
2074 nargs = 0;
2075 } else {
2076 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002077 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002078 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2079 return NULL;
2080 }
2081 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002082 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002083
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002084 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002085 if (indices == NULL) {
2086 PyErr_NoMemory();
2087 goto error;
2088 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002090 pools = PyTuple_New(npools);
2091 if (pools == NULL)
2092 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002093
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002094 for (i=0; i < nargs ; ++i) {
2095 PyObject *item = PyTuple_GET_ITEM(args, i);
2096 PyObject *pool = PySequence_Tuple(item);
2097 if (pool == NULL)
2098 goto error;
2099 PyTuple_SET_ITEM(pools, i, pool);
2100 indices[i] = 0;
2101 }
2102 for ( ; i < npools; ++i) {
2103 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2104 Py_INCREF(pool);
2105 PyTuple_SET_ITEM(pools, i, pool);
2106 indices[i] = 0;
2107 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002108
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002109 /* create productobject structure */
2110 lz = (productobject *)type->tp_alloc(type, 0);
2111 if (lz == NULL)
2112 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002114 lz->pools = pools;
2115 lz->indices = indices;
2116 lz->result = NULL;
2117 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002118
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002119 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002120
2121error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002122 if (indices != NULL)
2123 PyMem_Free(indices);
2124 Py_XDECREF(pools);
2125 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002126}
2127
2128static void
2129product_dealloc(productobject *lz)
2130{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002131 PyObject_GC_UnTrack(lz);
2132 Py_XDECREF(lz->pools);
2133 Py_XDECREF(lz->result);
2134 if (lz->indices != NULL)
2135 PyMem_Free(lz->indices);
2136 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002137}
2138
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002139static PyObject *
2140product_sizeof(productobject *lz, void *unused)
2141{
2142 Py_ssize_t res;
2143
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002144 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002145 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2146 return PyLong_FromSsize_t(res);
2147}
2148
2149PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2150
Christian Heimesc3f30c42008-02-22 16:37:40 +00002151static int
2152product_traverse(productobject *lz, visitproc visit, void *arg)
2153{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002154 Py_VISIT(lz->pools);
2155 Py_VISIT(lz->result);
2156 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002157}
2158
2159static PyObject *
2160product_next(productobject *lz)
2161{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002162 PyObject *pool;
2163 PyObject *elem;
2164 PyObject *oldelem;
2165 PyObject *pools = lz->pools;
2166 PyObject *result = lz->result;
2167 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2168 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002169
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002170 if (lz->stopped)
2171 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002172
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002173 if (result == NULL) {
2174 /* On the first pass, return an initial tuple filled with the
2175 first element from each pool. */
2176 result = PyTuple_New(npools);
2177 if (result == NULL)
2178 goto empty;
2179 lz->result = result;
2180 for (i=0; i < npools; i++) {
2181 pool = PyTuple_GET_ITEM(pools, i);
2182 if (PyTuple_GET_SIZE(pool) == 0)
2183 goto empty;
2184 elem = PyTuple_GET_ITEM(pool, 0);
2185 Py_INCREF(elem);
2186 PyTuple_SET_ITEM(result, i, elem);
2187 }
2188 } else {
2189 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002190
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002191 /* Copy the previous result tuple or re-use it if available */
2192 if (Py_REFCNT(result) > 1) {
2193 PyObject *old_result = result;
2194 result = PyTuple_New(npools);
2195 if (result == NULL)
2196 goto empty;
2197 lz->result = result;
2198 for (i=0; i < npools; i++) {
2199 elem = PyTuple_GET_ITEM(old_result, i);
2200 Py_INCREF(elem);
2201 PyTuple_SET_ITEM(result, i, elem);
2202 }
2203 Py_DECREF(old_result);
2204 }
2205 /* Now, we've got the only copy so we can update it in-place */
2206 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002207
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002208 /* Update the pool indices right-to-left. Only advance to the
2209 next pool when the previous one rolls-over */
2210 for (i=npools-1 ; i >= 0 ; i--) {
2211 pool = PyTuple_GET_ITEM(pools, i);
2212 indices[i]++;
2213 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2214 /* Roll-over and advance to next pool */
2215 indices[i] = 0;
2216 elem = PyTuple_GET_ITEM(pool, 0);
2217 Py_INCREF(elem);
2218 oldelem = PyTuple_GET_ITEM(result, i);
2219 PyTuple_SET_ITEM(result, i, elem);
2220 Py_DECREF(oldelem);
2221 } else {
2222 /* No rollover. Just increment and stop here. */
2223 elem = PyTuple_GET_ITEM(pool, indices[i]);
2224 Py_INCREF(elem);
2225 oldelem = PyTuple_GET_ITEM(result, i);
2226 PyTuple_SET_ITEM(result, i, elem);
2227 Py_DECREF(oldelem);
2228 break;
2229 }
2230 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002231
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002232 /* If i is negative, then the indices have all rolled-over
2233 and we're done. */
2234 if (i < 0)
2235 goto empty;
2236 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002237
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002238 Py_INCREF(result);
2239 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002240
2241empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002242 lz->stopped = 1;
2243 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002244}
2245
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002246static PyObject *
2247product_reduce(productobject *lz)
2248{
2249 if (lz->stopped) {
2250 return Py_BuildValue("O(())", Py_TYPE(lz));
2251 } else if (lz->result == NULL) {
2252 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2253 } else {
2254 PyObject *indices;
2255 Py_ssize_t n, i;
2256
2257 /* we must pickle the indices use them for setstate, and
2258 * additionally indicate that the iterator has started
2259 */
2260 n = PyTuple_GET_SIZE(lz->pools);
2261 indices = PyTuple_New(n);
2262 if (indices == NULL)
2263 return NULL;
2264 for (i=0; i<n; i++){
2265 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2266 if (!index) {
2267 Py_DECREF(indices);
2268 return NULL;
2269 }
2270 PyTuple_SET_ITEM(indices, i, index);
2271 }
2272 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2273 }
2274}
2275
2276static PyObject *
2277product_setstate(productobject *lz, PyObject *state)
2278{
2279 PyObject *result;
2280 Py_ssize_t n, i;
2281
2282 n = PyTuple_GET_SIZE(lz->pools);
2283 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2284 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2285 return NULL;
2286 }
2287 for (i=0; i<n; i++)
2288 {
2289 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2290 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002291 PyObject* pool;
2292 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002293 if (index < 0 && PyErr_Occurred())
2294 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002295 pool = PyTuple_GET_ITEM(lz->pools, i);
2296 poolsize = PyTuple_GET_SIZE(pool);
2297 if (poolsize == 0) {
2298 lz->stopped = 1;
2299 Py_RETURN_NONE;
2300 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002301 /* clamp the index */
2302 if (index < 0)
2303 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002304 else if (index > poolsize-1)
2305 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002306 lz->indices[i] = index;
2307 }
2308
2309 result = PyTuple_New(n);
2310 if (!result)
2311 return NULL;
2312 for (i=0; i<n; i++) {
2313 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2314 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2315 Py_INCREF(element);
2316 PyTuple_SET_ITEM(result, i, element);
2317 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002318 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002319 Py_RETURN_NONE;
2320}
2321
2322static PyMethodDef product_methods[] = {
2323 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2324 reduce_doc},
2325 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2326 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002327 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2328 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002329 {NULL, NULL} /* sentinel */
2330};
2331
Christian Heimesc3f30c42008-02-22 16:37:40 +00002332PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002333"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002334\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002335Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002336For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2337The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2338cycle in a manner similar to an odometer (with the rightmost element changing\n\
2339on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002340To compute the product of an iterable with itself, specify the number\n\
2341of repetitions with the optional repeat keyword argument. For example,\n\
2342product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002343product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2344product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2345
2346static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002347 PyVarObject_HEAD_INIT(NULL, 0)
2348 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002349 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002350 0, /* tp_itemsize */
2351 /* methods */
2352 (destructor)product_dealloc, /* tp_dealloc */
2353 0, /* tp_print */
2354 0, /* tp_getattr */
2355 0, /* tp_setattr */
2356 0, /* tp_reserved */
2357 0, /* tp_repr */
2358 0, /* tp_as_number */
2359 0, /* tp_as_sequence */
2360 0, /* tp_as_mapping */
2361 0, /* tp_hash */
2362 0, /* tp_call */
2363 0, /* tp_str */
2364 PyObject_GenericGetAttr, /* tp_getattro */
2365 0, /* tp_setattro */
2366 0, /* tp_as_buffer */
2367 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2368 Py_TPFLAGS_BASETYPE, /* tp_flags */
2369 product_doc, /* tp_doc */
2370 (traverseproc)product_traverse, /* tp_traverse */
2371 0, /* tp_clear */
2372 0, /* tp_richcompare */
2373 0, /* tp_weaklistoffset */
2374 PyObject_SelfIter, /* tp_iter */
2375 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002376 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002377 0, /* tp_members */
2378 0, /* tp_getset */
2379 0, /* tp_base */
2380 0, /* tp_dict */
2381 0, /* tp_descr_get */
2382 0, /* tp_descr_set */
2383 0, /* tp_dictoffset */
2384 0, /* tp_init */
2385 0, /* tp_alloc */
2386 product_new, /* tp_new */
2387 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002388};
2389
2390
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002391/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002392
2393typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002394 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002395 PyObject *pool; /* input converted to a tuple */
2396 Py_ssize_t *indices; /* one index per result element */
2397 PyObject *result; /* most recently returned result tuple */
2398 Py_ssize_t r; /* size of result tuple */
2399 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002400} combinationsobject;
2401
2402static PyTypeObject combinations_type;
2403
2404static PyObject *
2405combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2406{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002407 combinationsobject *co;
2408 Py_ssize_t n;
2409 Py_ssize_t r;
2410 PyObject *pool = NULL;
2411 PyObject *iterable = NULL;
2412 Py_ssize_t *indices = NULL;
2413 Py_ssize_t i;
2414 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002416 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2417 &iterable, &r))
2418 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002419
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002420 pool = PySequence_Tuple(iterable);
2421 if (pool == NULL)
2422 goto error;
2423 n = PyTuple_GET_SIZE(pool);
2424 if (r < 0) {
2425 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2426 goto error;
2427 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002428
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002429 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002430 if (indices == NULL) {
2431 PyErr_NoMemory();
2432 goto error;
2433 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002434
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 for (i=0 ; i<r ; i++)
2436 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002437
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002438 /* create combinationsobject structure */
2439 co = (combinationsobject *)type->tp_alloc(type, 0);
2440 if (co == NULL)
2441 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002443 co->pool = pool;
2444 co->indices = indices;
2445 co->result = NULL;
2446 co->r = r;
2447 co->stopped = r > n ? 1 : 0;
2448
2449 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002450
2451error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002452 if (indices != NULL)
2453 PyMem_Free(indices);
2454 Py_XDECREF(pool);
2455 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002456}
2457
2458static void
2459combinations_dealloc(combinationsobject *co)
2460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002461 PyObject_GC_UnTrack(co);
2462 Py_XDECREF(co->pool);
2463 Py_XDECREF(co->result);
2464 if (co->indices != NULL)
2465 PyMem_Free(co->indices);
2466 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002467}
2468
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002469static PyObject *
2470combinations_sizeof(combinationsobject *co, void *unused)
2471{
2472 Py_ssize_t res;
2473
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002474 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002475 res += co->r * sizeof(Py_ssize_t);
2476 return PyLong_FromSsize_t(res);
2477}
2478
Christian Heimes380f7f22008-02-28 11:19:05 +00002479static int
2480combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002482 Py_VISIT(co->pool);
2483 Py_VISIT(co->result);
2484 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002485}
2486
2487static PyObject *
2488combinations_next(combinationsobject *co)
2489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002490 PyObject *elem;
2491 PyObject *oldelem;
2492 PyObject *pool = co->pool;
2493 Py_ssize_t *indices = co->indices;
2494 PyObject *result = co->result;
2495 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2496 Py_ssize_t r = co->r;
2497 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002498
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002499 if (co->stopped)
2500 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002501
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002502 if (result == NULL) {
2503 /* On the first pass, initialize result tuple using the indices */
2504 result = PyTuple_New(r);
2505 if (result == NULL)
2506 goto empty;
2507 co->result = result;
2508 for (i=0; i<r ; i++) {
2509 index = indices[i];
2510 elem = PyTuple_GET_ITEM(pool, index);
2511 Py_INCREF(elem);
2512 PyTuple_SET_ITEM(result, i, elem);
2513 }
2514 } else {
2515 /* Copy the previous result tuple or re-use it if available */
2516 if (Py_REFCNT(result) > 1) {
2517 PyObject *old_result = result;
2518 result = PyTuple_New(r);
2519 if (result == NULL)
2520 goto empty;
2521 co->result = result;
2522 for (i=0; i<r ; i++) {
2523 elem = PyTuple_GET_ITEM(old_result, i);
2524 Py_INCREF(elem);
2525 PyTuple_SET_ITEM(result, i, elem);
2526 }
2527 Py_DECREF(old_result);
2528 }
2529 /* Now, we've got the only copy so we can update it in-place
2530 * CPython's empty tuple is a singleton and cached in
2531 * PyTuple's freelist.
2532 */
2533 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002535 /* Scan indices right-to-left until finding one that is not
2536 at its maximum (i + n - r). */
2537 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2538 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002539
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002540 /* If i is negative, then the indices are all at
2541 their maximum value and we're done. */
2542 if (i < 0)
2543 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002544
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002545 /* Increment the current index which we know is not at its
2546 maximum. Then move back to the right setting each index
2547 to its lowest possible value (one higher than the index
2548 to its left -- this maintains the sort order invariant). */
2549 indices[i]++;
2550 for (j=i+1 ; j<r ; j++)
2551 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002552
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002553 /* Update the result tuple for the new indices
2554 starting with i, the leftmost index that changed */
2555 for ( ; i<r ; i++) {
2556 index = indices[i];
2557 elem = PyTuple_GET_ITEM(pool, index);
2558 Py_INCREF(elem);
2559 oldelem = PyTuple_GET_ITEM(result, i);
2560 PyTuple_SET_ITEM(result, i, elem);
2561 Py_DECREF(oldelem);
2562 }
2563 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002564
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002565 Py_INCREF(result);
2566 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002567
2568empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002569 co->stopped = 1;
2570 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002571}
2572
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002573static PyObject *
2574combinations_reduce(combinationsobject *lz)
2575{
2576 if (lz->result == NULL) {
2577 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2578 } else if (lz->stopped) {
2579 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2580 } else {
2581 PyObject *indices;
2582 Py_ssize_t i;
2583
2584 /* we must pickle the indices and use them for setstate */
2585 indices = PyTuple_New(lz->r);
2586 if (!indices)
2587 return NULL;
2588 for (i=0; i<lz->r; i++)
2589 {
2590 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2591 if (!index) {
2592 Py_DECREF(indices);
2593 return NULL;
2594 }
2595 PyTuple_SET_ITEM(indices, i, index);
2596 }
2597
2598 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2599 }
2600}
2601
2602static PyObject *
2603combinations_setstate(combinationsobject *lz, PyObject *state)
2604{
2605 PyObject *result;
2606 Py_ssize_t i;
2607 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2608
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002609 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002610 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2611 return NULL;
2612 }
2613
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002614 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002615 Py_ssize_t max;
2616 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2617 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002618
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002619 if (index == -1 && PyErr_Occurred())
2620 return NULL; /* not an integer */
2621 max = i + n - lz->r;
2622 /* clamp the index (beware of negative max) */
2623 if (index > max)
2624 index = max;
2625 if (index < 0)
2626 index = 0;
2627 lz->indices[i] = index;
2628 }
2629
2630 result = PyTuple_New(lz->r);
2631 if (result == NULL)
2632 return NULL;
2633 for (i=0; i<lz->r; i++) {
2634 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2635 Py_INCREF(element);
2636 PyTuple_SET_ITEM(result, i, element);
2637 }
2638
Serhiy Storchaka48842712016-04-06 09:45:48 +03002639 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002640 Py_RETURN_NONE;
2641}
2642
2643static PyMethodDef combinations_methods[] = {
2644 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2645 reduce_doc},
2646 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2647 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002648 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2649 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002650 {NULL, NULL} /* sentinel */
2651};
2652
Christian Heimes380f7f22008-02-28 11:19:05 +00002653PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002654"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002655\n\
2656Return successive r-length combinations of elements in the iterable.\n\n\
2657combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2658
2659static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002660 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002661 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002662 sizeof(combinationsobject), /* tp_basicsize */
2663 0, /* tp_itemsize */
2664 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002665 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002666 0, /* tp_print */
2667 0, /* tp_getattr */
2668 0, /* tp_setattr */
2669 0, /* tp_reserved */
2670 0, /* tp_repr */
2671 0, /* tp_as_number */
2672 0, /* tp_as_sequence */
2673 0, /* tp_as_mapping */
2674 0, /* tp_hash */
2675 0, /* tp_call */
2676 0, /* tp_str */
2677 PyObject_GenericGetAttr, /* tp_getattro */
2678 0, /* tp_setattro */
2679 0, /* tp_as_buffer */
2680 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2681 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002682 combinations_doc, /* tp_doc */
2683 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002684 0, /* tp_clear */
2685 0, /* tp_richcompare */
2686 0, /* tp_weaklistoffset */
2687 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002688 (iternextfunc)combinations_next, /* tp_iternext */
2689 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002690 0, /* tp_members */
2691 0, /* tp_getset */
2692 0, /* tp_base */
2693 0, /* tp_dict */
2694 0, /* tp_descr_get */
2695 0, /* tp_descr_set */
2696 0, /* tp_dictoffset */
2697 0, /* tp_init */
2698 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002699 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002700 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002701};
2702
2703
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002704/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002705
2706/* Equivalent to:
2707
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 def combinations_with_replacement(iterable, r):
2709 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2710 # number items returned: (n+r-1)! / r! / (n-1)!
2711 pool = tuple(iterable)
2712 n = len(pool)
2713 indices = [0] * r
2714 yield tuple(pool[i] for i in indices)
2715 while 1:
2716 for i in reversed(range(r)):
2717 if indices[i] != n - 1:
2718 break
2719 else:
2720 return
2721 indices[i:] = [indices[i] + 1] * (r - i)
2722 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002723
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002724 def combinations_with_replacement2(iterable, r):
2725 'Alternate version that filters from product()'
2726 pool = tuple(iterable)
2727 n = len(pool)
2728 for indices in product(range(n), repeat=r):
2729 if sorted(indices) == list(indices):
2730 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002731*/
2732typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002733 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002734 PyObject *pool; /* input converted to a tuple */
2735 Py_ssize_t *indices; /* one index per result element */
2736 PyObject *result; /* most recently returned result tuple */
2737 Py_ssize_t r; /* size of result tuple */
2738 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002739} cwrobject;
2740
2741static PyTypeObject cwr_type;
2742
2743static PyObject *
2744cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2745{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002746 cwrobject *co;
2747 Py_ssize_t n;
2748 Py_ssize_t r;
2749 PyObject *pool = NULL;
2750 PyObject *iterable = NULL;
2751 Py_ssize_t *indices = NULL;
2752 Py_ssize_t i;
2753 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002754
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002755 if (!PyArg_ParseTupleAndKeywords(args, kwds,
2756 "On:combinations_with_replacement",
2757 kwargs, &iterable, &r))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002758 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 pool = PySequence_Tuple(iterable);
2761 if (pool == NULL)
2762 goto error;
2763 n = PyTuple_GET_SIZE(pool);
2764 if (r < 0) {
2765 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2766 goto error;
2767 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002768
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002769 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002770 if (indices == NULL) {
2771 PyErr_NoMemory();
2772 goto error;
2773 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002774
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002775 for (i=0 ; i<r ; i++)
2776 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002777
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002778 /* create cwrobject structure */
2779 co = (cwrobject *)type->tp_alloc(type, 0);
2780 if (co == NULL)
2781 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002782
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002783 co->pool = pool;
2784 co->indices = indices;
2785 co->result = NULL;
2786 co->r = r;
2787 co->stopped = !n && r;
2788
2789 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002790
2791error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002792 if (indices != NULL)
2793 PyMem_Free(indices);
2794 Py_XDECREF(pool);
2795 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002796}
2797
2798static void
2799cwr_dealloc(cwrobject *co)
2800{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002801 PyObject_GC_UnTrack(co);
2802 Py_XDECREF(co->pool);
2803 Py_XDECREF(co->result);
2804 if (co->indices != NULL)
2805 PyMem_Free(co->indices);
2806 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002807}
2808
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002809static PyObject *
2810cwr_sizeof(cwrobject *co, void *unused)
2811{
2812 Py_ssize_t res;
2813
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002814 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002815 res += co->r * sizeof(Py_ssize_t);
2816 return PyLong_FromSsize_t(res);
2817}
2818
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002819static int
2820cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2821{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 Py_VISIT(co->pool);
2823 Py_VISIT(co->result);
2824 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002825}
2826
2827static PyObject *
2828cwr_next(cwrobject *co)
2829{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 PyObject *elem;
2831 PyObject *oldelem;
2832 PyObject *pool = co->pool;
2833 Py_ssize_t *indices = co->indices;
2834 PyObject *result = co->result;
2835 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2836 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002837 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002838
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002839 if (co->stopped)
2840 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002841
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002842 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002843 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002844 result = PyTuple_New(r);
2845 if (result == NULL)
2846 goto empty;
2847 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002848 if (n > 0) {
2849 elem = PyTuple_GET_ITEM(pool, 0);
2850 for (i=0; i<r ; i++) {
2851 assert(indices[i] == 0);
2852 Py_INCREF(elem);
2853 PyTuple_SET_ITEM(result, i, elem);
2854 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002855 }
2856 } else {
2857 /* Copy the previous result tuple or re-use it if available */
2858 if (Py_REFCNT(result) > 1) {
2859 PyObject *old_result = result;
2860 result = PyTuple_New(r);
2861 if (result == NULL)
2862 goto empty;
2863 co->result = result;
2864 for (i=0; i<r ; i++) {
2865 elem = PyTuple_GET_ITEM(old_result, i);
2866 Py_INCREF(elem);
2867 PyTuple_SET_ITEM(result, i, elem);
2868 }
2869 Py_DECREF(old_result);
2870 }
2871 /* Now, we've got the only copy so we can update it in-place CPython's
2872 empty tuple is a singleton and cached in PyTuple's freelist. */
2873 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002874
Tim Peters9edb1682013-09-03 11:49:31 -05002875 /* Scan indices right-to-left until finding one that is not
2876 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002877 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2878 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002880 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002881 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002882 if (i < 0)
2883 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002885 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002886 maximum. Then set all to the right to the same value. */
2887 index = indices[i] + 1;
2888 assert(index < n);
2889 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002890 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002891 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002892 Py_INCREF(elem);
2893 oldelem = PyTuple_GET_ITEM(result, i);
2894 PyTuple_SET_ITEM(result, i, elem);
2895 Py_DECREF(oldelem);
2896 }
2897 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002899 Py_INCREF(result);
2900 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002901
2902empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002903 co->stopped = 1;
2904 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002905}
2906
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002907static PyObject *
2908cwr_reduce(cwrobject *lz)
2909{
2910 if (lz->result == NULL) {
2911 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2912 } else if (lz->stopped) {
2913 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2914 } else {
2915 PyObject *indices;
2916 Py_ssize_t i;
2917
2918 /* we must pickle the indices and use them for setstate */
2919 indices = PyTuple_New(lz->r);
2920 if (!indices)
2921 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002922 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002923 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2924 if (!index) {
2925 Py_DECREF(indices);
2926 return NULL;
2927 }
2928 PyTuple_SET_ITEM(indices, i, index);
2929 }
2930
2931 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2932 }
2933}
2934
2935static PyObject *
2936cwr_setstate(cwrobject *lz, PyObject *state)
2937{
2938 PyObject *result;
2939 Py_ssize_t n, i;
2940
2941 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2942 {
2943 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2944 return NULL;
2945 }
2946
2947 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002948 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002949 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2950 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002951
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002952 if (index < 0 && PyErr_Occurred())
2953 return NULL; /* not an integer */
2954 /* clamp the index */
2955 if (index < 0)
2956 index = 0;
2957 else if (index > n-1)
2958 index = n-1;
2959 lz->indices[i] = index;
2960 }
2961 result = PyTuple_New(lz->r);
2962 if (result == NULL)
2963 return NULL;
2964 for (i=0; i<lz->r; i++) {
2965 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2966 Py_INCREF(element);
2967 PyTuple_SET_ITEM(result, i, element);
2968 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002969 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002970 Py_RETURN_NONE;
2971}
2972
2973static PyMethodDef cwr_methods[] = {
2974 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2975 reduce_doc},
2976 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2977 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002978 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2979 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002980 {NULL, NULL} /* sentinel */
2981};
2982
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002983PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002984"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002985\n\
2986Return successive r-length combinations of elements in the iterable\n\
2987allowing individual elements to have successive repeats.\n\
2988combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2989
2990static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002991 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002992 "itertools.combinations_with_replacement", /* tp_name */
2993 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002994 0, /* tp_itemsize */
2995 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002996 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002997 0, /* tp_print */
2998 0, /* tp_getattr */
2999 0, /* tp_setattr */
3000 0, /* tp_reserved */
3001 0, /* tp_repr */
3002 0, /* tp_as_number */
3003 0, /* tp_as_sequence */
3004 0, /* tp_as_mapping */
3005 0, /* tp_hash */
3006 0, /* tp_call */
3007 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003008 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003009 0, /* tp_setattro */
3010 0, /* tp_as_buffer */
3011 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003012 Py_TPFLAGS_BASETYPE, /* tp_flags */
3013 cwr_doc, /* tp_doc */
3014 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003015 0, /* tp_clear */
3016 0, /* tp_richcompare */
3017 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003018 PyObject_SelfIter, /* tp_iter */
3019 (iternextfunc)cwr_next, /* tp_iternext */
3020 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003021 0, /* tp_members */
3022 0, /* tp_getset */
3023 0, /* tp_base */
3024 0, /* tp_dict */
3025 0, /* tp_descr_get */
3026 0, /* tp_descr_set */
3027 0, /* tp_dictoffset */
3028 0, /* tp_init */
3029 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003030 cwr_new, /* tp_new */
3031 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003032};
3033
3034
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003035/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003036
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003037def permutations(iterable, r=None):
3038 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3039 pool = tuple(iterable)
3040 n = len(pool)
3041 r = n if r is None else r
3042 indices = range(n)
3043 cycles = range(n-r+1, n+1)[::-1]
3044 yield tuple(pool[i] for i in indices[:r])
3045 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003046 for i in reversed(range(r)):
3047 cycles[i] -= 1
3048 if cycles[i] == 0:
3049 indices[i:] = indices[i+1:] + indices[i:i+1]
3050 cycles[i] = n - i
3051 else:
3052 j = cycles[i]
3053 indices[i], indices[-j] = indices[-j], indices[i]
3054 yield tuple(pool[i] for i in indices[:r])
3055 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003056 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003057 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003058*/
3059
3060typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003062 PyObject *pool; /* input converted to a tuple */
3063 Py_ssize_t *indices; /* one index per element in the pool */
3064 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3065 PyObject *result; /* most recently returned result tuple */
3066 Py_ssize_t r; /* size of result tuple */
3067 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003068} permutationsobject;
3069
3070static PyTypeObject permutations_type;
3071
3072static PyObject *
3073permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3074{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003075 permutationsobject *po;
3076 Py_ssize_t n;
3077 Py_ssize_t r;
3078 PyObject *robj = Py_None;
3079 PyObject *pool = NULL;
3080 PyObject *iterable = NULL;
3081 Py_ssize_t *indices = NULL;
3082 Py_ssize_t *cycles = NULL;
3083 Py_ssize_t i;
3084 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003085
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003086 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3087 &iterable, &robj))
3088 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003090 pool = PySequence_Tuple(iterable);
3091 if (pool == NULL)
3092 goto error;
3093 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003094
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003095 r = n;
3096 if (robj != Py_None) {
3097 if (!PyLong_Check(robj)) {
3098 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3099 goto error;
3100 }
3101 r = PyLong_AsSsize_t(robj);
3102 if (r == -1 && PyErr_Occurred())
3103 goto error;
3104 }
3105 if (r < 0) {
3106 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3107 goto error;
3108 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003109
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003110 indices = PyMem_New(Py_ssize_t, n);
3111 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003112 if (indices == NULL || cycles == NULL) {
3113 PyErr_NoMemory();
3114 goto error;
3115 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003117 for (i=0 ; i<n ; i++)
3118 indices[i] = i;
3119 for (i=0 ; i<r ; i++)
3120 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 /* create permutationsobject structure */
3123 po = (permutationsobject *)type->tp_alloc(type, 0);
3124 if (po == NULL)
3125 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003127 po->pool = pool;
3128 po->indices = indices;
3129 po->cycles = cycles;
3130 po->result = NULL;
3131 po->r = r;
3132 po->stopped = r > n ? 1 : 0;
3133
3134 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003135
3136error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003137 if (indices != NULL)
3138 PyMem_Free(indices);
3139 if (cycles != NULL)
3140 PyMem_Free(cycles);
3141 Py_XDECREF(pool);
3142 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003143}
3144
3145static void
3146permutations_dealloc(permutationsobject *po)
3147{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003148 PyObject_GC_UnTrack(po);
3149 Py_XDECREF(po->pool);
3150 Py_XDECREF(po->result);
3151 PyMem_Free(po->indices);
3152 PyMem_Free(po->cycles);
3153 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003154}
3155
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003156static PyObject *
3157permutations_sizeof(permutationsobject *po, void *unused)
3158{
3159 Py_ssize_t res;
3160
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003161 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003162 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3163 res += po->r * sizeof(Py_ssize_t);
3164 return PyLong_FromSsize_t(res);
3165}
3166
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003167static int
3168permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3169{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003170 Py_VISIT(po->pool);
3171 Py_VISIT(po->result);
3172 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003173}
3174
3175static PyObject *
3176permutations_next(permutationsobject *po)
3177{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003178 PyObject *elem;
3179 PyObject *oldelem;
3180 PyObject *pool = po->pool;
3181 Py_ssize_t *indices = po->indices;
3182 Py_ssize_t *cycles = po->cycles;
3183 PyObject *result = po->result;
3184 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3185 Py_ssize_t r = po->r;
3186 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003187
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003188 if (po->stopped)
3189 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003190
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003191 if (result == NULL) {
3192 /* On the first pass, initialize result tuple using the indices */
3193 result = PyTuple_New(r);
3194 if (result == NULL)
3195 goto empty;
3196 po->result = result;
3197 for (i=0; i<r ; i++) {
3198 index = indices[i];
3199 elem = PyTuple_GET_ITEM(pool, index);
3200 Py_INCREF(elem);
3201 PyTuple_SET_ITEM(result, i, elem);
3202 }
3203 } else {
3204 if (n == 0)
3205 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003206
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003207 /* Copy the previous result tuple or re-use it if available */
3208 if (Py_REFCNT(result) > 1) {
3209 PyObject *old_result = result;
3210 result = PyTuple_New(r);
3211 if (result == NULL)
3212 goto empty;
3213 po->result = result;
3214 for (i=0; i<r ; i++) {
3215 elem = PyTuple_GET_ITEM(old_result, i);
3216 Py_INCREF(elem);
3217 PyTuple_SET_ITEM(result, i, elem);
3218 }
3219 Py_DECREF(old_result);
3220 }
3221 /* Now, we've got the only copy so we can update it in-place */
3222 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003223
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003224 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3225 for (i=r-1 ; i>=0 ; i--) {
3226 cycles[i] -= 1;
3227 if (cycles[i] == 0) {
3228 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3229 index = indices[i];
3230 for (j=i ; j<n-1 ; j++)
3231 indices[j] = indices[j+1];
3232 indices[n-1] = index;
3233 cycles[i] = n - i;
3234 } else {
3235 j = cycles[i];
3236 index = indices[i];
3237 indices[i] = indices[n-j];
3238 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003239
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003240 for (k=i; k<r ; k++) {
3241 /* start with i, the leftmost element that changed */
3242 /* yield tuple(pool[k] for k in indices[:r]) */
3243 index = indices[k];
3244 elem = PyTuple_GET_ITEM(pool, index);
3245 Py_INCREF(elem);
3246 oldelem = PyTuple_GET_ITEM(result, k);
3247 PyTuple_SET_ITEM(result, k, elem);
3248 Py_DECREF(oldelem);
3249 }
3250 break;
3251 }
3252 }
3253 /* If i is negative, then the cycles have all
3254 rolled-over and we're done. */
3255 if (i < 0)
3256 goto empty;
3257 }
3258 Py_INCREF(result);
3259 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003260
3261empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003262 po->stopped = 1;
3263 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003264}
3265
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003266static PyObject *
3267permutations_reduce(permutationsobject *po)
3268{
3269 if (po->result == NULL) {
3270 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3271 } else if (po->stopped) {
3272 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3273 } else {
3274 PyObject *indices=NULL, *cycles=NULL;
3275 Py_ssize_t n, i;
3276
3277 /* we must pickle the indices and cycles and use them for setstate */
3278 n = PyTuple_GET_SIZE(po->pool);
3279 indices = PyTuple_New(n);
3280 if (indices == NULL)
3281 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003282 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003283 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3284 if (!index)
3285 goto err;
3286 PyTuple_SET_ITEM(indices, i, index);
3287 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003288
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003289 cycles = PyTuple_New(po->r);
3290 if (cycles == NULL)
3291 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003292 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003293 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3294 if (!index)
3295 goto err;
3296 PyTuple_SET_ITEM(cycles, i, index);
3297 }
3298 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3299 po->pool, po->r,
3300 indices, cycles);
3301 err:
3302 Py_XDECREF(indices);
3303 Py_XDECREF(cycles);
3304 return NULL;
3305 }
3306}
3307
3308static PyObject *
3309permutations_setstate(permutationsobject *po, PyObject *state)
3310{
3311 PyObject *indices, *cycles, *result;
3312 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003313
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003314 if (!PyTuple_Check(state)) {
3315 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3316 return NULL;
3317 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003318 if (!PyArg_ParseTuple(state, "O!O!",
3319 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003320 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003321 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003322 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003323
3324 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003325 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003326 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3327 return NULL;
3328 }
3329
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003330 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003331 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3332 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3333 if (index < 0 && PyErr_Occurred())
3334 return NULL; /* not an integer */
3335 /* clamp the index */
3336 if (index < 0)
3337 index = 0;
3338 else if (index > n-1)
3339 index = n-1;
3340 po->indices[i] = index;
3341 }
3342
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* indexObject = PyTuple_GET_ITEM(cycles, i);
3345 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3346 if (index < 0 && PyErr_Occurred())
3347 return NULL; /* not an integer */
3348 if (index < 1)
3349 index = 1;
3350 else if (index > n-i)
3351 index = n-i;
3352 po->cycles[i] = index;
3353 }
3354 result = PyTuple_New(po->r);
3355 if (result == NULL)
3356 return NULL;
3357 for (i=0; i<po->r; i++) {
3358 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3359 Py_INCREF(element);
3360 PyTuple_SET_ITEM(result, i, element);
3361 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003362 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003363 Py_RETURN_NONE;
3364}
3365
3366static PyMethodDef permuations_methods[] = {
3367 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3368 reduce_doc},
3369 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3370 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003371 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3372 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003373 {NULL, NULL} /* sentinel */
3374};
3375
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003376PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003377"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003378\n\
3379Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003380permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003381
3382static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003383 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003384 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003385 sizeof(permutationsobject), /* tp_basicsize */
3386 0, /* tp_itemsize */
3387 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003388 (destructor)permutations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003389 0, /* tp_print */
3390 0, /* tp_getattr */
3391 0, /* tp_setattr */
3392 0, /* tp_reserved */
3393 0, /* tp_repr */
3394 0, /* tp_as_number */
3395 0, /* tp_as_sequence */
3396 0, /* tp_as_mapping */
3397 0, /* tp_hash */
3398 0, /* tp_call */
3399 0, /* tp_str */
3400 PyObject_GenericGetAttr, /* tp_getattro */
3401 0, /* tp_setattro */
3402 0, /* tp_as_buffer */
3403 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3404 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003405 permutations_doc, /* tp_doc */
3406 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003407 0, /* tp_clear */
3408 0, /* tp_richcompare */
3409 0, /* tp_weaklistoffset */
3410 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003411 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003412 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003413 0, /* tp_members */
3414 0, /* tp_getset */
3415 0, /* tp_base */
3416 0, /* tp_dict */
3417 0, /* tp_descr_get */
3418 0, /* tp_descr_set */
3419 0, /* tp_dictoffset */
3420 0, /* tp_init */
3421 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003422 permutations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003423 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003424};
3425
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003426/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003427
3428typedef struct {
3429 PyObject_HEAD
3430 PyObject *total;
3431 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003432 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003433} accumulateobject;
3434
3435static PyTypeObject accumulate_type;
3436
3437static PyObject *
3438accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3439{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003440 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003441 PyObject *iterable;
3442 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003443 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003444 accumulateobject *lz;
3445
Raymond Hettinger5d446132011-03-27 18:52:10 -07003446 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3447 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003448 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003449
3450 /* Get iterator. */
3451 it = PyObject_GetIter(iterable);
3452 if (it == NULL)
3453 return NULL;
3454
Raymond Hettinger482ba772010-12-01 22:48:00 +00003455 /* create accumulateobject structure */
3456 lz = (accumulateobject *)type->tp_alloc(type, 0);
3457 if (lz == NULL) {
3458 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003459 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003460 }
3461
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003462 if (binop != Py_None) {
3463 Py_XINCREF(binop);
3464 lz->binop = binop;
3465 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003466 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003467 lz->it = it;
3468 return (PyObject *)lz;
3469}
3470
3471static void
3472accumulate_dealloc(accumulateobject *lz)
3473{
3474 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003475 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003476 Py_XDECREF(lz->total);
3477 Py_XDECREF(lz->it);
3478 Py_TYPE(lz)->tp_free(lz);
3479}
3480
3481static int
3482accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3483{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003484 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003485 Py_VISIT(lz->it);
3486 Py_VISIT(lz->total);
3487 return 0;
3488}
3489
3490static PyObject *
3491accumulate_next(accumulateobject *lz)
3492{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003493 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003494
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003495 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003496 if (val == NULL)
3497 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003498
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003499 if (lz->total == NULL) {
3500 Py_INCREF(val);
3501 lz->total = val;
3502 return lz->total;
3503 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003504
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003505 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003506 newtotal = PyNumber_Add(lz->total, val);
3507 else
3508 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003509 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003510 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003511 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003512
Raymond Hettinger482ba772010-12-01 22:48:00 +00003513 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003514 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003515 return newtotal;
3516}
3517
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003518static PyObject *
3519accumulate_reduce(accumulateobject *lz)
3520{
Serhiy Storchakad5516252016-03-06 14:00:45 +02003521 if (lz->total == Py_None) {
3522 PyObject *it;
3523
3524 if (PyType_Ready(&chain_type) < 0)
3525 return NULL;
3526 if (PyType_Ready(&islice_type) < 0)
3527 return NULL;
3528 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3529 lz->total, lz->it);
3530 if (it == NULL)
3531 return NULL;
3532 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3533 it, lz->binop ? lz->binop : Py_None);
3534 if (it == NULL)
3535 return NULL;
3536 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3537 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003538 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3539 lz->it, lz->binop?lz->binop:Py_None,
3540 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003541}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003542
3543static PyObject *
3544accumulate_setstate(accumulateobject *lz, PyObject *state)
3545{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003546 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003547 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003548 Py_RETURN_NONE;
3549}
3550
3551static PyMethodDef accumulate_methods[] = {
3552 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3553 reduce_doc},
3554 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3555 setstate_doc},
3556 {NULL, NULL} /* sentinel */
3557};
3558
Raymond Hettinger482ba772010-12-01 22:48:00 +00003559PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003560"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003561\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003562Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003563
3564static PyTypeObject accumulate_type = {
3565 PyVarObject_HEAD_INIT(NULL, 0)
3566 "itertools.accumulate", /* tp_name */
3567 sizeof(accumulateobject), /* tp_basicsize */
3568 0, /* tp_itemsize */
3569 /* methods */
3570 (destructor)accumulate_dealloc, /* tp_dealloc */
3571 0, /* tp_print */
3572 0, /* tp_getattr */
3573 0, /* tp_setattr */
3574 0, /* tp_reserved */
3575 0, /* tp_repr */
3576 0, /* tp_as_number */
3577 0, /* tp_as_sequence */
3578 0, /* tp_as_mapping */
3579 0, /* tp_hash */
3580 0, /* tp_call */
3581 0, /* tp_str */
3582 PyObject_GenericGetAttr, /* tp_getattro */
3583 0, /* tp_setattro */
3584 0, /* tp_as_buffer */
3585 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3586 Py_TPFLAGS_BASETYPE, /* tp_flags */
3587 accumulate_doc, /* tp_doc */
3588 (traverseproc)accumulate_traverse, /* tp_traverse */
3589 0, /* tp_clear */
3590 0, /* tp_richcompare */
3591 0, /* tp_weaklistoffset */
3592 PyObject_SelfIter, /* tp_iter */
3593 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003594 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003595 0, /* tp_members */
3596 0, /* tp_getset */
3597 0, /* tp_base */
3598 0, /* tp_dict */
3599 0, /* tp_descr_get */
3600 0, /* tp_descr_set */
3601 0, /* tp_dictoffset */
3602 0, /* tp_init */
3603 0, /* tp_alloc */
3604 accumulate_new, /* tp_new */
3605 PyObject_GC_Del, /* tp_free */
3606};
3607
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003608
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003609/* compress object ************************************************************/
3610
3611/* Equivalent to:
3612
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003613 def compress(data, selectors):
3614 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3615 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003616*/
3617
3618typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003619 PyObject_HEAD
3620 PyObject *data;
3621 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003622} compressobject;
3623
3624static PyTypeObject compress_type;
3625
3626static PyObject *
3627compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3628{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003629 PyObject *seq1, *seq2;
3630 PyObject *data=NULL, *selectors=NULL;
3631 compressobject *lz;
3632 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003633
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003634 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3635 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003636
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003637 data = PyObject_GetIter(seq1);
3638 if (data == NULL)
3639 goto fail;
3640 selectors = PyObject_GetIter(seq2);
3641 if (selectors == NULL)
3642 goto fail;
3643
3644 /* create compressobject structure */
3645 lz = (compressobject *)type->tp_alloc(type, 0);
3646 if (lz == NULL)
3647 goto fail;
3648 lz->data = data;
3649 lz->selectors = selectors;
3650 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003651
3652fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003653 Py_XDECREF(data);
3654 Py_XDECREF(selectors);
3655 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003656}
3657
3658static void
3659compress_dealloc(compressobject *lz)
3660{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003661 PyObject_GC_UnTrack(lz);
3662 Py_XDECREF(lz->data);
3663 Py_XDECREF(lz->selectors);
3664 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003665}
3666
3667static int
3668compress_traverse(compressobject *lz, visitproc visit, void *arg)
3669{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003670 Py_VISIT(lz->data);
3671 Py_VISIT(lz->selectors);
3672 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003673}
3674
3675static PyObject *
3676compress_next(compressobject *lz)
3677{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003678 PyObject *data = lz->data, *selectors = lz->selectors;
3679 PyObject *datum, *selector;
3680 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3681 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3682 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003683
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003684 while (1) {
3685 /* Steps: get datum, get selector, evaluate selector.
3686 Order is important (to match the pure python version
3687 in terms of which input gets a chance to raise an
3688 exception first).
3689 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003690
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003691 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003692 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003693 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003694
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003695 selector = selectornext(selectors);
3696 if (selector == NULL) {
3697 Py_DECREF(datum);
3698 return NULL;
3699 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003700
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003701 ok = PyObject_IsTrue(selector);
3702 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003703 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 return datum;
3705 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003706 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003707 return NULL;
3708 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003709}
3710
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003711static PyObject *
3712compress_reduce(compressobject *lz)
3713{
3714 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3715 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003716}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003717
3718static PyMethodDef compress_methods[] = {
3719 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3720 reduce_doc},
3721 {NULL, NULL} /* sentinel */
3722};
3723
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003724PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003725"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003726\n\
3727Return data elements corresponding to true selector elements.\n\
3728Forms a shorter iterator from selected data elements using the\n\
3729selectors to choose the data elements.");
3730
3731static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003732 PyVarObject_HEAD_INIT(NULL, 0)
3733 "itertools.compress", /* tp_name */
3734 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003735 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003736 /* methods */
3737 (destructor)compress_dealloc, /* tp_dealloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003738 0, /* tp_print */
3739 0, /* tp_getattr */
3740 0, /* tp_setattr */
3741 0, /* tp_reserved */
3742 0, /* tp_repr */
3743 0, /* tp_as_number */
3744 0, /* tp_as_sequence */
3745 0, /* tp_as_mapping */
3746 0, /* tp_hash */
3747 0, /* tp_call */
3748 0, /* tp_str */
3749 PyObject_GenericGetAttr, /* tp_getattro */
3750 0, /* tp_setattro */
3751 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003752 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003753 Py_TPFLAGS_BASETYPE, /* tp_flags */
3754 compress_doc, /* tp_doc */
3755 (traverseproc)compress_traverse, /* tp_traverse */
3756 0, /* tp_clear */
3757 0, /* tp_richcompare */
3758 0, /* tp_weaklistoffset */
3759 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003760 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003761 compress_methods, /* tp_methods */
3762 0, /* tp_members */
3763 0, /* tp_getset */
3764 0, /* tp_base */
3765 0, /* tp_dict */
3766 0, /* tp_descr_get */
3767 0, /* tp_descr_set */
3768 0, /* tp_dictoffset */
3769 0, /* tp_init */
3770 0, /* tp_alloc */
3771 compress_new, /* tp_new */
3772 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003773};
3774
3775
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003776/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003777
3778typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003779 PyObject_HEAD
3780 PyObject *func;
3781 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003782} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003783
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003784static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003785
3786static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003787filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003788{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003789 PyObject *func, *seq;
3790 PyObject *it;
3791 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003792
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003793 if (type == &filterfalse_type &&
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03003794 !_PyArg_NoKeywords("filterfalse", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003795 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003796
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003797 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3798 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003799
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003800 /* Get iterator. */
3801 it = PyObject_GetIter(seq);
3802 if (it == NULL)
3803 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003804
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003805 /* create filterfalseobject structure */
3806 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3807 if (lz == NULL) {
3808 Py_DECREF(it);
3809 return NULL;
3810 }
3811 Py_INCREF(func);
3812 lz->func = func;
3813 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003814
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003815 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003816}
3817
3818static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003819filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003820{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003821 PyObject_GC_UnTrack(lz);
3822 Py_XDECREF(lz->func);
3823 Py_XDECREF(lz->it);
3824 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003825}
3826
3827static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003828filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003829{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003830 Py_VISIT(lz->it);
3831 Py_VISIT(lz->func);
3832 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003833}
3834
3835static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003836filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003837{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003838 PyObject *item;
3839 PyObject *it = lz->it;
3840 long ok;
3841 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003842
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003843 iternext = *Py_TYPE(it)->tp_iternext;
3844 for (;;) {
3845 item = iternext(it);
3846 if (item == NULL)
3847 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003848
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003849 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3850 ok = PyObject_IsTrue(item);
3851 } else {
3852 PyObject *good;
Victor Stinnerde4ae3d2016-12-04 22:59:09 +01003853 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003854 if (good == NULL) {
3855 Py_DECREF(item);
3856 return NULL;
3857 }
3858 ok = PyObject_IsTrue(good);
3859 Py_DECREF(good);
3860 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003861 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003862 return item;
3863 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003864 if (ok < 0)
3865 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003866 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003867}
3868
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003869static PyObject *
3870filterfalse_reduce(filterfalseobject *lz)
3871{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003872 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3873}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003874
3875static PyMethodDef filterfalse_methods[] = {
3876 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3877 reduce_doc},
3878 {NULL, NULL} /* sentinel */
3879};
3880
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003881PyDoc_STRVAR(filterfalse_doc,
3882"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003883\n\
3884Return those items of sequence for which function(item) is false.\n\
3885If function is None, return the items that are false.");
3886
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003887static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003888 PyVarObject_HEAD_INIT(NULL, 0)
3889 "itertools.filterfalse", /* tp_name */
3890 sizeof(filterfalseobject), /* tp_basicsize */
3891 0, /* tp_itemsize */
3892 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003893 (destructor)filterfalse_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003894 0, /* tp_print */
3895 0, /* tp_getattr */
3896 0, /* tp_setattr */
3897 0, /* tp_reserved */
3898 0, /* tp_repr */
3899 0, /* tp_as_number */
3900 0, /* tp_as_sequence */
3901 0, /* tp_as_mapping */
3902 0, /* tp_hash */
3903 0, /* tp_call */
3904 0, /* tp_str */
3905 PyObject_GenericGetAttr, /* tp_getattro */
3906 0, /* tp_setattro */
3907 0, /* tp_as_buffer */
3908 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3909 Py_TPFLAGS_BASETYPE, /* tp_flags */
3910 filterfalse_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003911 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003912 0, /* tp_clear */
3913 0, /* tp_richcompare */
3914 0, /* tp_weaklistoffset */
3915 PyObject_SelfIter, /* tp_iter */
3916 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003917 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003918 0, /* tp_members */
3919 0, /* tp_getset */
3920 0, /* tp_base */
3921 0, /* tp_dict */
3922 0, /* tp_descr_get */
3923 0, /* tp_descr_set */
3924 0, /* tp_dictoffset */
3925 0, /* tp_init */
3926 0, /* tp_alloc */
3927 filterfalse_new, /* tp_new */
3928 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003929};
3930
3931
3932/* count object ************************************************************/
3933
3934typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 PyObject_HEAD
3936 Py_ssize_t cnt;
3937 PyObject *long_cnt;
3938 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003939} countobject;
3940
Raymond Hettinger30729212009-02-12 06:28:27 +00003941/* Counting logic and invariants:
3942
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003943fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003944
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003945 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003946 Advances with: cnt += 1
3947 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003948
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003949slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003950
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003951 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3952 All counting is done with python objects (no overflows or underflows).
3953 Advances with: long_cnt += long_step
3954 Step may be zero -- effectively a slow version of repeat(cnt).
3955 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003956*/
3957
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003958static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003959
3960static PyObject *
3961count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3962{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003963 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003964 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003965 Py_ssize_t cnt = 0;
3966 PyObject *long_cnt = NULL;
3967 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003968 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003969 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003970
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003971 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3972 kwlist, &long_cnt, &long_step))
3973 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003974
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3976 (long_step != NULL && !PyNumber_Check(long_step))) {
3977 PyErr_SetString(PyExc_TypeError, "a number is required");
3978 return NULL;
3979 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003980
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003981 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
3982 (long_step == NULL || PyLong_Check(long_step));
3983
3984 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003985 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03003986 if (fast_mode) {
3987 assert(PyLong_Check(long_cnt));
3988 cnt = PyLong_AsSsize_t(long_cnt);
3989 if (cnt == -1 && PyErr_Occurred()) {
3990 PyErr_Clear();
3991 fast_mode = 0;
3992 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003993 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003994 } else {
3995 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003996 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003997 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03003998 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003999
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004000 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004001 if (long_step == NULL)
4002 long_step = _PyLong_One;
4003 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004005 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004007 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004008 if (fast_mode) {
4009 assert(PyLong_Check(long_step));
4010 step = PyLong_AsLong(long_step);
4011 if (step != 1) {
4012 fast_mode = 0;
4013 if (step == -1 && PyErr_Occurred())
4014 PyErr_Clear();
4015 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004016 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004017
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004018 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004019 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004020 else
4021 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004022
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004023 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4024 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4025 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004026 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004027
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004028 /* create countobject structure */
4029 lz = (countobject *)type->tp_alloc(type, 0);
4030 if (lz == NULL) {
4031 Py_XDECREF(long_cnt);
4032 return NULL;
4033 }
4034 lz->cnt = cnt;
4035 lz->long_cnt = long_cnt;
4036 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004037
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004038 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004039}
4040
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004041static void
4042count_dealloc(countobject *lz)
4043{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004044 PyObject_GC_UnTrack(lz);
4045 Py_XDECREF(lz->long_cnt);
4046 Py_XDECREF(lz->long_step);
4047 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004048}
4049
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004050static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004051count_traverse(countobject *lz, visitproc visit, void *arg)
4052{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004053 Py_VISIT(lz->long_cnt);
4054 Py_VISIT(lz->long_step);
4055 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004056}
4057
4058static PyObject *
4059count_nextlong(countobject *lz)
4060{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004061 PyObject *long_cnt;
4062 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004063
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004064 long_cnt = lz->long_cnt;
4065 if (long_cnt == NULL) {
4066 /* Switch to slow_mode */
4067 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4068 if (long_cnt == NULL)
4069 return NULL;
4070 }
4071 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004073 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4074 if (stepped_up == NULL)
4075 return NULL;
4076 lz->long_cnt = stepped_up;
4077 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004078}
4079
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004080static PyObject *
4081count_next(countobject *lz)
4082{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004083 if (lz->cnt == PY_SSIZE_T_MAX)
4084 return count_nextlong(lz);
4085 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004086}
4087
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004088static PyObject *
4089count_repr(countobject *lz)
4090{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004091 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004092 return PyUnicode_FromFormat("%s(%zd)",
4093 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004094
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004095 if (PyLong_Check(lz->long_step)) {
4096 long step = PyLong_AsLong(lz->long_step);
4097 if (step == -1 && PyErr_Occurred()) {
4098 PyErr_Clear();
4099 }
4100 if (step == 1) {
4101 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004102 return PyUnicode_FromFormat("%s(%R)",
4103 _PyType_Name(Py_TYPE(lz)),
4104 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 }
4106 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004107 return PyUnicode_FromFormat("%s(%R, %R)",
4108 _PyType_Name(Py_TYPE(lz)),
4109 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004110}
4111
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004112static PyObject *
4113count_reduce(countobject *lz)
4114{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004115 if (lz->cnt == PY_SSIZE_T_MAX)
4116 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4117 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004118}
4119
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004120static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004121 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004122 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004124};
4125
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004126PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004127 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004128\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004129Return a count object whose .__next__() method returns consecutive values.\n\
4130Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004131 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004132 x = firstval\n\
4133 while 1:\n\
4134 yield x\n\
4135 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004136
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004137static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004138 PyVarObject_HEAD_INIT(NULL, 0)
4139 "itertools.count", /* tp_name */
4140 sizeof(countobject), /* tp_basicsize */
4141 0, /* tp_itemsize */
4142 /* methods */
4143 (destructor)count_dealloc, /* tp_dealloc */
4144 0, /* tp_print */
4145 0, /* tp_getattr */
4146 0, /* tp_setattr */
4147 0, /* tp_reserved */
4148 (reprfunc)count_repr, /* tp_repr */
4149 0, /* tp_as_number */
4150 0, /* tp_as_sequence */
4151 0, /* tp_as_mapping */
4152 0, /* tp_hash */
4153 0, /* tp_call */
4154 0, /* tp_str */
4155 PyObject_GenericGetAttr, /* tp_getattro */
4156 0, /* tp_setattro */
4157 0, /* tp_as_buffer */
4158 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004159 Py_TPFLAGS_BASETYPE, /* tp_flags */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004160 count_doc, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004161 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004162 0, /* tp_clear */
4163 0, /* tp_richcompare */
4164 0, /* tp_weaklistoffset */
4165 PyObject_SelfIter, /* tp_iter */
4166 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004167 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004168 0, /* tp_members */
4169 0, /* tp_getset */
4170 0, /* tp_base */
4171 0, /* tp_dict */
4172 0, /* tp_descr_get */
4173 0, /* tp_descr_set */
4174 0, /* tp_dictoffset */
4175 0, /* tp_init */
4176 0, /* tp_alloc */
4177 count_new, /* tp_new */
4178 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004179};
4180
4181
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004182/* repeat object ************************************************************/
4183
4184typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004185 PyObject_HEAD
4186 PyObject *element;
4187 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004188} repeatobject;
4189
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004190static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004191
4192static PyObject *
4193repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4194{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004195 repeatobject *ro;
4196 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004197 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004198 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004199
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004200 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4201 &element, &cnt))
4202 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004203
Raymond Hettinger97d35552014-06-24 21:36:58 -07004204 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004205 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004206 /* Does user supply times argument? */
4207 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004208 cnt = 0;
4209
4210 ro = (repeatobject *)type->tp_alloc(type, 0);
4211 if (ro == NULL)
4212 return NULL;
4213 Py_INCREF(element);
4214 ro->element = element;
4215 ro->cnt = cnt;
4216 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004217}
4218
4219static void
4220repeat_dealloc(repeatobject *ro)
4221{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004222 PyObject_GC_UnTrack(ro);
4223 Py_XDECREF(ro->element);
4224 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004225}
4226
4227static int
4228repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4229{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004230 Py_VISIT(ro->element);
4231 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004232}
4233
4234static PyObject *
4235repeat_next(repeatobject *ro)
4236{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004237 if (ro->cnt == 0)
4238 return NULL;
4239 if (ro->cnt > 0)
4240 ro->cnt--;
4241 Py_INCREF(ro->element);
4242 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004243}
4244
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004245static PyObject *
4246repeat_repr(repeatobject *ro)
4247{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004248 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004249 return PyUnicode_FromFormat("%s(%R)",
4250 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004251 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004252 return PyUnicode_FromFormat("%s(%R, %zd)",
4253 _PyType_Name(Py_TYPE(ro)), ro->element,
4254 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004255}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004256
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004257static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004258repeat_len(repeatobject *ro)
4259{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004260 if (ro->cnt == -1) {
4261 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4262 return NULL;
4263 }
4264 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004265}
4266
Armin Rigof5b3e362006-02-11 21:32:43 +00004267PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004268
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004269static PyObject *
4270repeat_reduce(repeatobject *ro)
4271{
4272 /* unpickle this so that a new repeat iterator is constructed with an
4273 * object, then call __setstate__ on it to set cnt
4274 */
4275 if (ro->cnt >= 0)
4276 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4277 else
4278 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4279}
4280
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004281static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004282 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004283 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004284 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004285};
4286
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004287PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004288"repeat(object [,times]) -> create an iterator which returns the object\n\
4289for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004290endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004291
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004292static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004293 PyVarObject_HEAD_INIT(NULL, 0)
4294 "itertools.repeat", /* tp_name */
4295 sizeof(repeatobject), /* tp_basicsize */
4296 0, /* tp_itemsize */
4297 /* methods */
4298 (destructor)repeat_dealloc, /* tp_dealloc */
4299 0, /* tp_print */
4300 0, /* tp_getattr */
4301 0, /* tp_setattr */
4302 0, /* tp_reserved */
4303 (reprfunc)repeat_repr, /* tp_repr */
4304 0, /* tp_as_number */
4305 0, /* tp_as_sequence */
4306 0, /* tp_as_mapping */
4307 0, /* tp_hash */
4308 0, /* tp_call */
4309 0, /* tp_str */
4310 PyObject_GenericGetAttr, /* tp_getattro */
4311 0, /* tp_setattro */
4312 0, /* tp_as_buffer */
4313 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4314 Py_TPFLAGS_BASETYPE, /* tp_flags */
4315 repeat_doc, /* tp_doc */
4316 (traverseproc)repeat_traverse, /* tp_traverse */
4317 0, /* tp_clear */
4318 0, /* tp_richcompare */
4319 0, /* tp_weaklistoffset */
4320 PyObject_SelfIter, /* tp_iter */
4321 (iternextfunc)repeat_next, /* tp_iternext */
4322 repeat_methods, /* tp_methods */
4323 0, /* tp_members */
4324 0, /* tp_getset */
4325 0, /* tp_base */
4326 0, /* tp_dict */
4327 0, /* tp_descr_get */
4328 0, /* tp_descr_set */
4329 0, /* tp_dictoffset */
4330 0, /* tp_init */
4331 0, /* tp_alloc */
4332 repeat_new, /* tp_new */
4333 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004334};
4335
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004336/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004337
4338typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004339 PyObject_HEAD
4340 Py_ssize_t tuplesize;
4341 Py_ssize_t numactive;
4342 PyObject *ittuple; /* tuple of iterators */
4343 PyObject *result;
4344 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004345} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004346
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004347static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004348
4349static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004350zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004351{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004352 ziplongestobject *lz;
4353 Py_ssize_t i;
4354 PyObject *ittuple; /* tuple of iterators */
4355 PyObject *result;
4356 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004357 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004358
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004359 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004360 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004361 if (fillvalue == NULL || PyDict_GET_SIZE(kwds) > 1) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004362 PyErr_SetString(PyExc_TypeError,
4363 "zip_longest() got an unexpected keyword argument");
4364 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004365 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004366 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004368 /* args must be a tuple */
4369 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004370 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004372 /* obtain iterators */
4373 ittuple = PyTuple_New(tuplesize);
4374 if (ittuple == NULL)
4375 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004376 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004377 PyObject *item = PyTuple_GET_ITEM(args, i);
4378 PyObject *it = PyObject_GetIter(item);
4379 if (it == NULL) {
4380 if (PyErr_ExceptionMatches(PyExc_TypeError))
4381 PyErr_Format(PyExc_TypeError,
4382 "zip_longest argument #%zd must support iteration",
4383 i+1);
4384 Py_DECREF(ittuple);
4385 return NULL;
4386 }
4387 PyTuple_SET_ITEM(ittuple, i, it);
4388 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004389
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004390 /* create a result holder */
4391 result = PyTuple_New(tuplesize);
4392 if (result == NULL) {
4393 Py_DECREF(ittuple);
4394 return NULL;
4395 }
4396 for (i=0 ; i < tuplesize ; i++) {
4397 Py_INCREF(Py_None);
4398 PyTuple_SET_ITEM(result, i, Py_None);
4399 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004400
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004401 /* create ziplongestobject structure */
4402 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4403 if (lz == NULL) {
4404 Py_DECREF(ittuple);
4405 Py_DECREF(result);
4406 return NULL;
4407 }
4408 lz->ittuple = ittuple;
4409 lz->tuplesize = tuplesize;
4410 lz->numactive = tuplesize;
4411 lz->result = result;
4412 Py_INCREF(fillvalue);
4413 lz->fillvalue = fillvalue;
4414 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004415}
4416
4417static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004418zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004419{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004420 PyObject_GC_UnTrack(lz);
4421 Py_XDECREF(lz->ittuple);
4422 Py_XDECREF(lz->result);
4423 Py_XDECREF(lz->fillvalue);
4424 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004425}
4426
4427static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004428zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004429{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004430 Py_VISIT(lz->ittuple);
4431 Py_VISIT(lz->result);
4432 Py_VISIT(lz->fillvalue);
4433 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004434}
4435
4436static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004437zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004438{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004439 Py_ssize_t i;
4440 Py_ssize_t tuplesize = lz->tuplesize;
4441 PyObject *result = lz->result;
4442 PyObject *it;
4443 PyObject *item;
4444 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004446 if (tuplesize == 0)
4447 return NULL;
4448 if (lz->numactive == 0)
4449 return NULL;
4450 if (Py_REFCNT(result) == 1) {
4451 Py_INCREF(result);
4452 for (i=0 ; i < tuplesize ; i++) {
4453 it = PyTuple_GET_ITEM(lz->ittuple, i);
4454 if (it == NULL) {
4455 Py_INCREF(lz->fillvalue);
4456 item = lz->fillvalue;
4457 } else {
4458 item = PyIter_Next(it);
4459 if (item == NULL) {
4460 lz->numactive -= 1;
4461 if (lz->numactive == 0 || PyErr_Occurred()) {
4462 lz->numactive = 0;
4463 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004464 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004465 } else {
4466 Py_INCREF(lz->fillvalue);
4467 item = lz->fillvalue;
4468 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4469 Py_DECREF(it);
4470 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004471 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004472 }
4473 olditem = PyTuple_GET_ITEM(result, i);
4474 PyTuple_SET_ITEM(result, i, item);
4475 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004476 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004477 } else {
4478 result = PyTuple_New(tuplesize);
4479 if (result == NULL)
4480 return NULL;
4481 for (i=0 ; i < tuplesize ; i++) {
4482 it = PyTuple_GET_ITEM(lz->ittuple, i);
4483 if (it == NULL) {
4484 Py_INCREF(lz->fillvalue);
4485 item = lz->fillvalue;
4486 } else {
4487 item = PyIter_Next(it);
4488 if (item == NULL) {
4489 lz->numactive -= 1;
4490 if (lz->numactive == 0 || PyErr_Occurred()) {
4491 lz->numactive = 0;
4492 Py_DECREF(result);
4493 return NULL;
4494 } else {
4495 Py_INCREF(lz->fillvalue);
4496 item = lz->fillvalue;
4497 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4498 Py_DECREF(it);
4499 }
4500 }
4501 }
4502 PyTuple_SET_ITEM(result, i, item);
4503 }
4504 }
4505 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004506}
4507
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004508static PyObject *
4509zip_longest_reduce(ziplongestobject *lz)
4510{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004511
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004512 /* Create a new tuple with empty sequences where appropriate to pickle.
4513 * Then use setstate to set the fillvalue
4514 */
4515 int i;
4516 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004517
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004518 if (args == NULL)
4519 return NULL;
4520 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4521 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4522 if (elem == NULL) {
4523 elem = PyTuple_New(0);
4524 if (elem == NULL) {
4525 Py_DECREF(args);
4526 return NULL;
4527 }
4528 } else
4529 Py_INCREF(elem);
4530 PyTuple_SET_ITEM(args, i, elem);
4531 }
4532 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4533}
4534
4535static PyObject *
4536zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4537{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004538 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004539 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004540 Py_RETURN_NONE;
4541}
4542
4543static PyMethodDef zip_longest_methods[] = {
4544 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4545 reduce_doc},
4546 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4547 setstate_doc},
4548 {NULL, NULL} /* sentinel */
4549};
4550
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004551PyDoc_STRVAR(zip_longest_doc,
4552"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004553\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004554Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004555the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004556method continues until the longest iterable in the argument sequence\n\
4557is exhausted and then it raises StopIteration. When the shorter iterables\n\
4558are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4559defaults to None or can be specified by a keyword argument.\n\
4560");
4561
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004562static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004563 PyVarObject_HEAD_INIT(NULL, 0)
4564 "itertools.zip_longest", /* tp_name */
4565 sizeof(ziplongestobject), /* tp_basicsize */
4566 0, /* tp_itemsize */
4567 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004568 (destructor)zip_longest_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004569 0, /* tp_print */
4570 0, /* tp_getattr */
4571 0, /* tp_setattr */
4572 0, /* tp_reserved */
4573 0, /* tp_repr */
4574 0, /* tp_as_number */
4575 0, /* tp_as_sequence */
4576 0, /* tp_as_mapping */
4577 0, /* tp_hash */
4578 0, /* tp_call */
4579 0, /* tp_str */
4580 PyObject_GenericGetAttr, /* tp_getattro */
4581 0, /* tp_setattro */
4582 0, /* tp_as_buffer */
4583 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4584 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004585 zip_longest_doc, /* tp_doc */
4586 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004587 0, /* tp_clear */
4588 0, /* tp_richcompare */
4589 0, /* tp_weaklistoffset */
4590 PyObject_SelfIter, /* tp_iter */
4591 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004592 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004593 0, /* tp_members */
4594 0, /* tp_getset */
4595 0, /* tp_base */
4596 0, /* tp_dict */
4597 0, /* tp_descr_get */
4598 0, /* tp_descr_set */
4599 0, /* tp_dictoffset */
4600 0, /* tp_init */
4601 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004602 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004603 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004604};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004605
4606/* module level code ********************************************************/
4607
4608PyDoc_STRVAR(module_doc,
4609"Functional tools for creating and using iterators.\n\
4610\n\
4611Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004612count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004613cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004614repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004615\n\
4616Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004617accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004618chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004619chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004620compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4621dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4622groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004623filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004624islice(seq, [start,] stop [, step]) --> elements from\n\
4625 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004626starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004627tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004628takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004629zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004630\n\
4631Combinatoric generators:\n\
4632product(p, q, ... [repeat=1]) --> cartesian product\n\
4633permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004634combinations(p, r)\n\
4635combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004636");
4637
4638
Raymond Hettingerad983e72003-11-12 14:32:26 +00004639static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004640 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4641 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004642};
4643
Martin v. Löwis1a214512008-06-11 05:26:20 +00004644
4645static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004646 PyModuleDef_HEAD_INIT,
4647 "itertools",
4648 module_doc,
4649 -1,
4650 module_methods,
4651 NULL,
4652 NULL,
4653 NULL,
4654 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004655};
4656
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004657PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004658PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004659{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004660 int i;
4661 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004662 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004663 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004664 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004665 &combinations_type,
4666 &cwr_type,
4667 &cycle_type,
4668 &dropwhile_type,
4669 &takewhile_type,
4670 &islice_type,
4671 &starmap_type,
4672 &chain_type,
4673 &compress_type,
4674 &filterfalse_type,
4675 &count_type,
4676 &ziplongest_type,
4677 &permutations_type,
4678 &product_type,
4679 &repeat_type,
4680 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004681 &_grouper_type,
4682 &tee_type,
4683 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004684 NULL
4685 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004686
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004687 Py_TYPE(&teedataobject_type) = &PyType_Type;
4688 m = PyModule_Create(&itertoolsmodule);
4689 if (m == NULL)
4690 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004691
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004692 for (i=0 ; typelist[i] != NULL ; i++) {
4693 if (PyType_Ready(typelist[i]) < 0)
4694 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004695 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004696 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004697 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004698 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004699
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004700 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004701}