blob: cac02ad75a8f509e3c6211771483b307fb5db8ba [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007*/
8
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00009
10/* groupby object ***********************************************************/
11
12typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000013 PyObject_HEAD
14 PyObject *it;
15 PyObject *keyfunc;
16 PyObject *tgtkey;
17 PyObject *currkey;
18 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000019} groupbyobject;
20
21static PyTypeObject groupby_type;
22static PyObject *_grouper_create(groupbyobject *, PyObject *);
23
24static PyObject *
25groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
26{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000027 static char *kwargs[] = {"iterable", "key", NULL};
28 groupbyobject *gbo;
29 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000031 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
32 &it, &keyfunc))
33 return NULL;
34
35 gbo = (groupbyobject *)type->tp_alloc(type, 0);
36 if (gbo == NULL)
37 return NULL;
38 gbo->tgtkey = NULL;
39 gbo->currkey = NULL;
40 gbo->currvalue = NULL;
41 gbo->keyfunc = keyfunc;
42 Py_INCREF(keyfunc);
43 gbo->it = PyObject_GetIter(it);
44 if (gbo->it == NULL) {
45 Py_DECREF(gbo);
46 return NULL;
47 }
48 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000049}
50
51static void
52groupby_dealloc(groupbyobject *gbo)
53{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000054 PyObject_GC_UnTrack(gbo);
55 Py_XDECREF(gbo->it);
56 Py_XDECREF(gbo->keyfunc);
57 Py_XDECREF(gbo->tgtkey);
58 Py_XDECREF(gbo->currkey);
59 Py_XDECREF(gbo->currvalue);
60 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000061}
62
63static int
64groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
65{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000066 Py_VISIT(gbo->it);
67 Py_VISIT(gbo->keyfunc);
68 Py_VISIT(gbo->tgtkey);
69 Py_VISIT(gbo->currkey);
70 Py_VISIT(gbo->currvalue);
71 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000072}
73
74static PyObject *
75groupby_next(groupbyobject *gbo)
76{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000077 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000078
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 /* skip to next iteration group */
80 for (;;) {
81 if (gbo->currkey == NULL)
82 /* pass */;
83 else if (gbo->tgtkey == NULL)
84 break;
85 else {
86 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000088 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
89 gbo->currkey, Py_EQ);
90 if (rcmp == -1)
91 return NULL;
92 else if (rcmp == 0)
93 break;
94 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000096 newvalue = PyIter_Next(gbo->it);
97 if (newvalue == NULL)
98 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000100 if (gbo->keyfunc == Py_None) {
101 newkey = newvalue;
102 Py_INCREF(newvalue);
103 } else {
104 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
105 newvalue, NULL);
106 if (newkey == NULL) {
107 Py_DECREF(newvalue);
108 return NULL;
109 }
110 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000112 tmp = gbo->currkey;
113 gbo->currkey = newkey;
114 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000115
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000116 tmp = gbo->currvalue;
117 gbo->currvalue = newvalue;
118 Py_XDECREF(tmp);
119 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000121 Py_INCREF(gbo->currkey);
122 tmp = gbo->tgtkey;
123 gbo->tgtkey = gbo->currkey;
124 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000125
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000126 grouper = _grouper_create(gbo, gbo->tgtkey);
127 if (grouper == NULL)
128 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000130 r = PyTuple_Pack(2, gbo->currkey, grouper);
131 Py_DECREF(grouper);
132 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000133}
134
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000135static PyObject *
136groupby_reduce(groupbyobject *lz)
137{
138 /* reduce as a 'new' call with an optional 'setstate' if groupby
139 * has started
140 */
141 PyObject *value;
142 if (lz->tgtkey && lz->currkey && lz->currvalue)
143 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
144 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
145 else
146 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
147 lz->it, lz->keyfunc);
148
149 return value;
150}
151
152PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
153
154static PyObject *
155groupby_setstate(groupbyobject *lz, PyObject *state)
156{
157 PyObject *currkey, *currvalue, *tgtkey;
158 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
159 return NULL;
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200160 Py_INCREF(currkey);
161 Py_SETREF(lz->currkey, currkey);
162 Py_INCREF(currvalue);
163 Py_SETREF(lz->currvalue, currvalue);
164 Py_INCREF(tgtkey);
165 Py_SETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000166 Py_RETURN_NONE;
167}
168
169PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
170
171static PyMethodDef groupby_methods[] = {
172 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
173 reduce_doc},
174 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
175 setstate_doc},
176 {NULL, NULL} /* sentinel */
177};
178
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000179PyDoc_STRVAR(groupby_doc,
180"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
181(key, sub-iterator) grouped by each value of key(value).\n");
182
183static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000184 PyVarObject_HEAD_INIT(NULL, 0)
185 "itertools.groupby", /* tp_name */
186 sizeof(groupbyobject), /* tp_basicsize */
187 0, /* tp_itemsize */
188 /* methods */
189 (destructor)groupby_dealloc, /* tp_dealloc */
190 0, /* tp_print */
191 0, /* tp_getattr */
192 0, /* tp_setattr */
193 0, /* tp_reserved */
194 0, /* tp_repr */
195 0, /* tp_as_number */
196 0, /* tp_as_sequence */
197 0, /* tp_as_mapping */
198 0, /* tp_hash */
199 0, /* tp_call */
200 0, /* tp_str */
201 PyObject_GenericGetAttr, /* tp_getattro */
202 0, /* tp_setattro */
203 0, /* tp_as_buffer */
204 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
205 Py_TPFLAGS_BASETYPE, /* tp_flags */
206 groupby_doc, /* tp_doc */
207 (traverseproc)groupby_traverse, /* tp_traverse */
208 0, /* tp_clear */
209 0, /* tp_richcompare */
210 0, /* tp_weaklistoffset */
211 PyObject_SelfIter, /* tp_iter */
212 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000213 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000214 0, /* tp_members */
215 0, /* tp_getset */
216 0, /* tp_base */
217 0, /* tp_dict */
218 0, /* tp_descr_get */
219 0, /* tp_descr_set */
220 0, /* tp_dictoffset */
221 0, /* tp_init */
222 0, /* tp_alloc */
223 groupby_new, /* tp_new */
224 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000225};
226
227
228/* _grouper object (internal) ************************************************/
229
230typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000231 PyObject_HEAD
232 PyObject *parent;
233 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000234} _grouperobject;
235
236static PyTypeObject _grouper_type;
237
238static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000239_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
240{
241 PyObject *parent, *tgtkey;
242
243 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
244 return NULL;
245
246 return _grouper_create((groupbyobject*) parent, tgtkey);
247}
248
249static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000250_grouper_create(groupbyobject *parent, PyObject *tgtkey)
251{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000252 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000254 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
255 if (igo == NULL)
256 return NULL;
257 igo->parent = (PyObject *)parent;
258 Py_INCREF(parent);
259 igo->tgtkey = tgtkey;
260 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000261
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000262 PyObject_GC_Track(igo);
263 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264}
265
266static void
267_grouper_dealloc(_grouperobject *igo)
268{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000269 PyObject_GC_UnTrack(igo);
270 Py_DECREF(igo->parent);
271 Py_DECREF(igo->tgtkey);
272 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000273}
274
275static int
276_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
277{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000278 Py_VISIT(igo->parent);
279 Py_VISIT(igo->tgtkey);
280 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000281}
282
283static PyObject *
284_grouper_next(_grouperobject *igo)
285{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000286 groupbyobject *gbo = (groupbyobject *)igo->parent;
287 PyObject *newvalue, *newkey, *r;
288 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000290 if (gbo->currvalue == NULL) {
291 newvalue = PyIter_Next(gbo->it);
292 if (newvalue == NULL)
293 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (gbo->keyfunc == Py_None) {
296 newkey = newvalue;
297 Py_INCREF(newvalue);
298 } else {
299 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
300 newvalue, NULL);
301 if (newkey == NULL) {
302 Py_DECREF(newvalue);
303 return NULL;
304 }
305 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000306
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000307 assert(gbo->currkey == NULL);
308 gbo->currkey = newkey;
309 gbo->currvalue = newvalue;
310 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 assert(gbo->currkey != NULL);
313 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
314 if (rcmp <= 0)
315 /* got any error or current group is end */
316 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000317
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000318 r = gbo->currvalue;
319 gbo->currvalue = NULL;
320 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000321
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000322 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000323}
324
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000325static PyObject *
326_grouper_reduce(_grouperobject *lz)
327{
328 return Py_BuildValue("O(OO)", Py_TYPE(lz),
329 lz->parent, lz->tgtkey);
330}
331
332static PyMethodDef _grouper_methods[] = {
333 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
334 reduce_doc},
335 {NULL, NULL} /* sentinel */
336};
337
338
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000339static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000340 PyVarObject_HEAD_INIT(NULL, 0)
341 "itertools._grouper", /* tp_name */
342 sizeof(_grouperobject), /* tp_basicsize */
343 0, /* tp_itemsize */
344 /* methods */
345 (destructor)_grouper_dealloc, /* tp_dealloc */
346 0, /* tp_print */
347 0, /* tp_getattr */
348 0, /* tp_setattr */
349 0, /* tp_reserved */
350 0, /* tp_repr */
351 0, /* tp_as_number */
352 0, /* tp_as_sequence */
353 0, /* tp_as_mapping */
354 0, /* tp_hash */
355 0, /* tp_call */
356 0, /* tp_str */
357 PyObject_GenericGetAttr, /* tp_getattro */
358 0, /* tp_setattro */
359 0, /* tp_as_buffer */
360 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
361 0, /* tp_doc */
362 (traverseproc)_grouper_traverse,/* tp_traverse */
363 0, /* tp_clear */
364 0, /* tp_richcompare */
365 0, /* tp_weaklistoffset */
366 PyObject_SelfIter, /* tp_iter */
367 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000368 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000369 0, /* tp_members */
370 0, /* tp_getset */
371 0, /* tp_base */
372 0, /* tp_dict */
373 0, /* tp_descr_get */
374 0, /* tp_descr_set */
375 0, /* tp_dictoffset */
376 0, /* tp_init */
377 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000378 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000379 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000380};
381
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000382
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000383
Raymond Hettingerad983e72003-11-12 14:32:26 +0000384/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000385
Raymond Hettingerad983e72003-11-12 14:32:26 +0000386/* The teedataobject pre-allocates space for LINKCELLS number of objects.
387 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000388 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 the other structure members including PyHEAD overhead. The larger the
390 value, the less memory overhead per object and the less time spent
391 allocating/deallocating new links. The smaller the number, the less
392 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000393*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000395
396typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000397 PyObject_HEAD
398 PyObject *it;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200399 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000400 PyObject *nextlink;
401 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000402} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000403
404typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000405 PyObject_HEAD
406 teedataobject *dataobj;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200407 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000408 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000409} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000410
Raymond Hettingerad983e72003-11-12 14:32:26 +0000411static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000412
413static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000414teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000416 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000418 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
419 if (tdo == NULL)
420 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000421
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000422 tdo->numread = 0;
423 tdo->nextlink = NULL;
424 Py_INCREF(it);
425 tdo->it = it;
426 PyObject_GC_Track(tdo);
427 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000428}
429
430static PyObject *
431teedataobject_jumplink(teedataobject *tdo)
432{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000433 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000434 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000435 Py_XINCREF(tdo->nextlink);
436 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000437}
438
439static PyObject *
440teedataobject_getitem(teedataobject *tdo, int i)
441{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000442 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000443
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000444 assert(i < LINKCELLS);
445 if (i < tdo->numread)
446 value = tdo->values[i];
447 else {
448 /* this is the lead iterator, so fetch more data */
449 assert(i == tdo->numread);
450 value = PyIter_Next(tdo->it);
451 if (value == NULL)
452 return NULL;
453 tdo->numread++;
454 tdo->values[i] = value;
455 }
456 Py_INCREF(value);
457 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000458}
459
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000460static int
461teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
462{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 int i;
464 Py_VISIT(tdo->it);
465 for (i = 0; i < tdo->numread; i++)
466 Py_VISIT(tdo->values[i]);
467 Py_VISIT(tdo->nextlink);
468 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000469}
470
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200471static void
472teedataobject_safe_decref(PyObject *obj)
473{
474 while (obj && Py_TYPE(obj) == &teedataobject_type &&
475 Py_REFCNT(obj) == 1) {
476 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
477 ((teedataobject *)obj)->nextlink = NULL;
478 Py_DECREF(obj);
479 obj = nextlink;
480 }
481 Py_XDECREF(obj);
482}
483
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000484static int
485teedataobject_clear(teedataobject *tdo)
486{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000487 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200488 PyObject *tmp;
489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 Py_CLEAR(tdo->it);
491 for (i=0 ; i<tdo->numread ; i++)
492 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200493 tmp = tdo->nextlink;
494 tdo->nextlink = NULL;
495 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000496 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000497}
498
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000499static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000501{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000502 PyObject_GC_UnTrack(tdo);
503 teedataobject_clear(tdo);
504 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000505}
506
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000507static PyObject *
508teedataobject_reduce(teedataobject *tdo)
509{
510 int i;
511 /* create a temporary list of already iterated values */
512 PyObject *values = PyList_New(tdo->numread);
513 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 = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000580 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 */
600 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
601 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 Storchaka5a57ade2015-12-24 10:35:59 +0200634 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;
739 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
740 return NULL;
741 if (index < 0 || index > LINKCELLS) {
742 PyErr_SetString(PyExc_ValueError, "Index out of range");
743 return NULL;
744 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200745 Py_INCREF(tdo);
746 Py_SETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000747 to->index = index;
748 Py_RETURN_NONE;
749}
750
Raymond Hettingerad983e72003-11-12 14:32:26 +0000751PyDoc_STRVAR(teeobject_doc,
752"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000753
Raymond Hettingerad983e72003-11-12 14:32:26 +0000754static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000755 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000756 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
757 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000758 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000759};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760
761static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000762 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000763 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000764 sizeof(teeobject), /* tp_basicsize */
765 0, /* tp_itemsize */
766 /* methods */
767 (destructor)tee_dealloc, /* tp_dealloc */
768 0, /* tp_print */
769 0, /* tp_getattr */
770 0, /* tp_setattr */
771 0, /* tp_reserved */
772 0, /* tp_repr */
773 0, /* tp_as_number */
774 0, /* tp_as_sequence */
775 0, /* tp_as_mapping */
776 0, /* tp_hash */
777 0, /* tp_call */
778 0, /* tp_str */
779 0, /* tp_getattro */
780 0, /* tp_setattro */
781 0, /* tp_as_buffer */
782 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
783 teeobject_doc, /* tp_doc */
784 (traverseproc)tee_traverse, /* tp_traverse */
785 (inquiry)tee_clear, /* tp_clear */
786 0, /* tp_richcompare */
787 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
788 PyObject_SelfIter, /* tp_iter */
789 (iternextfunc)tee_next, /* tp_iternext */
790 tee_methods, /* tp_methods */
791 0, /* tp_members */
792 0, /* tp_getset */
793 0, /* tp_base */
794 0, /* tp_dict */
795 0, /* tp_descr_get */
796 0, /* tp_descr_set */
797 0, /* tp_dictoffset */
798 0, /* tp_init */
799 0, /* tp_alloc */
800 tee_new, /* tp_new */
801 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000802};
803
Raymond Hettingerad983e72003-11-12 14:32:26 +0000804static PyObject *
805tee(PyObject *self, PyObject *args)
806{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000807 Py_ssize_t i, n=2;
808 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200809 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000810
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000811 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
812 return NULL;
813 if (n < 0) {
814 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
815 return NULL;
816 }
817 result = PyTuple_New(n);
818 if (result == NULL)
819 return NULL;
820 if (n == 0)
821 return result;
822 it = PyObject_GetIter(iterable);
823 if (it == NULL) {
824 Py_DECREF(result);
825 return NULL;
826 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200827 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000828 copyable = tee_fromiterable(it);
829 Py_DECREF(it);
830 if (copyable == NULL) {
831 Py_DECREF(result);
832 return NULL;
833 }
834 } else
835 copyable = it;
836 PyTuple_SET_ITEM(result, 0, copyable);
837 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200838
839 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000840 if (copyable == NULL) {
841 Py_DECREF(result);
842 return NULL;
843 }
844 PyTuple_SET_ITEM(result, i, copyable);
845 }
846 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000847}
848
849PyDoc_STRVAR(tee_doc,
850"tee(iterable, n=2) --> tuple of n independent iterators.");
851
852
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000853/* cycle object **********************************************************/
854
855typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000856 PyObject_HEAD
857 PyObject *it;
858 PyObject *saved;
859 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000860} cycleobject;
861
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000862static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000863
864static PyObject *
865cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
866{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000867 PyObject *it;
868 PyObject *iterable;
869 PyObject *saved;
870 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000871
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000872 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
873 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000875 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
876 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000878 /* Get iterator. */
879 it = PyObject_GetIter(iterable);
880 if (it == NULL)
881 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000882
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000883 saved = PyList_New(0);
884 if (saved == NULL) {
885 Py_DECREF(it);
886 return NULL;
887 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 /* create cycleobject structure */
890 lz = (cycleobject *)type->tp_alloc(type, 0);
891 if (lz == NULL) {
892 Py_DECREF(it);
893 Py_DECREF(saved);
894 return NULL;
895 }
896 lz->it = it;
897 lz->saved = saved;
898 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000899
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000900 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000901}
902
903static void
904cycle_dealloc(cycleobject *lz)
905{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000906 PyObject_GC_UnTrack(lz);
907 Py_XDECREF(lz->saved);
908 Py_XDECREF(lz->it);
909 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000910}
911
912static int
913cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
914{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000915 Py_VISIT(lz->it);
916 Py_VISIT(lz->saved);
917 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000918}
919
920static PyObject *
921cycle_next(cycleobject *lz)
922{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000923 PyObject *item;
924 PyObject *it;
925 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000927 while (1) {
928 item = PyIter_Next(lz->it);
929 if (item != NULL) {
930 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
931 Py_DECREF(item);
932 return NULL;
933 }
934 return item;
935 }
936 if (PyErr_Occurred()) {
937 if (PyErr_ExceptionMatches(PyExc_StopIteration))
938 PyErr_Clear();
939 else
940 return NULL;
941 }
942 if (PyList_Size(lz->saved) == 0)
943 return NULL;
944 it = PyObject_GetIter(lz->saved);
945 if (it == NULL)
946 return NULL;
947 tmp = lz->it;
948 lz->it = it;
949 lz->firstpass = 1;
950 Py_DECREF(tmp);
951 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000952}
953
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000954static PyObject *
955cycle_reduce(cycleobject *lz)
956{
957 /* Create a new cycle with the iterator tuple, then set
958 * the saved state on it.
959 */
Serhiy Storchakad269b5e2013-01-25 13:38:56 +0200960 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000961 lz->it, lz->saved, lz->firstpass);
962 }
963
964static PyObject *
965cycle_setstate(cycleobject *lz, PyObject *state)
966{
967 PyObject *saved=NULL;
968 int firstpass;
969 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
970 return NULL;
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200971 Py_XINCREF(saved);
972 Py_SETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000973 lz->firstpass = firstpass != 0;
974 Py_RETURN_NONE;
975}
976
977static PyMethodDef cycle_methods[] = {
978 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
979 reduce_doc},
980 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
981 setstate_doc},
982 {NULL, NULL} /* sentinel */
983};
984
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000985PyDoc_STRVAR(cycle_doc,
986"cycle(iterable) --> cycle object\n\
987\n\
988Return elements from the iterable until it is exhausted.\n\
989Then repeat the sequence indefinitely.");
990
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000991static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000992 PyVarObject_HEAD_INIT(NULL, 0)
993 "itertools.cycle", /* tp_name */
994 sizeof(cycleobject), /* tp_basicsize */
995 0, /* tp_itemsize */
996 /* methods */
997 (destructor)cycle_dealloc, /* tp_dealloc */
998 0, /* tp_print */
999 0, /* tp_getattr */
1000 0, /* tp_setattr */
1001 0, /* tp_reserved */
1002 0, /* tp_repr */
1003 0, /* tp_as_number */
1004 0, /* tp_as_sequence */
1005 0, /* tp_as_mapping */
1006 0, /* tp_hash */
1007 0, /* tp_call */
1008 0, /* tp_str */
1009 PyObject_GenericGetAttr, /* tp_getattro */
1010 0, /* tp_setattro */
1011 0, /* tp_as_buffer */
1012 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1013 Py_TPFLAGS_BASETYPE, /* tp_flags */
1014 cycle_doc, /* tp_doc */
1015 (traverseproc)cycle_traverse, /* tp_traverse */
1016 0, /* tp_clear */
1017 0, /* tp_richcompare */
1018 0, /* tp_weaklistoffset */
1019 PyObject_SelfIter, /* tp_iter */
1020 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001021 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001022 0, /* tp_members */
1023 0, /* tp_getset */
1024 0, /* tp_base */
1025 0, /* tp_dict */
1026 0, /* tp_descr_get */
1027 0, /* tp_descr_set */
1028 0, /* tp_dictoffset */
1029 0, /* tp_init */
1030 0, /* tp_alloc */
1031 cycle_new, /* tp_new */
1032 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001033};
1034
1035
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036/* dropwhile object **********************************************************/
1037
1038typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001039 PyObject_HEAD
1040 PyObject *func;
1041 PyObject *it;
1042 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001043} dropwhileobject;
1044
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001045static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001046
1047static PyObject *
1048dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1049{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001050 PyObject *func, *seq;
1051 PyObject *it;
1052 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001054 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1055 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001057 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1058 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001060 /* Get iterator. */
1061 it = PyObject_GetIter(seq);
1062 if (it == NULL)
1063 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 /* create dropwhileobject structure */
1066 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1067 if (lz == NULL) {
1068 Py_DECREF(it);
1069 return NULL;
1070 }
1071 Py_INCREF(func);
1072 lz->func = func;
1073 lz->it = it;
1074 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001076 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001077}
1078
1079static void
1080dropwhile_dealloc(dropwhileobject *lz)
1081{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001082 PyObject_GC_UnTrack(lz);
1083 Py_XDECREF(lz->func);
1084 Py_XDECREF(lz->it);
1085 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001086}
1087
1088static int
1089dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1090{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001091 Py_VISIT(lz->it);
1092 Py_VISIT(lz->func);
1093 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094}
1095
1096static PyObject *
1097dropwhile_next(dropwhileobject *lz)
1098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001099 PyObject *item, *good;
1100 PyObject *it = lz->it;
1101 long ok;
1102 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001103
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001104 iternext = *Py_TYPE(it)->tp_iternext;
1105 for (;;) {
1106 item = iternext(it);
1107 if (item == NULL)
1108 return NULL;
1109 if (lz->start == 1)
1110 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001112 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1113 if (good == NULL) {
1114 Py_DECREF(item);
1115 return NULL;
1116 }
1117 ok = PyObject_IsTrue(good);
1118 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001119 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 lz->start = 1;
1121 return item;
1122 }
1123 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001124 if (ok < 0)
1125 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001126 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001127}
1128
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001129static PyObject *
1130dropwhile_reduce(dropwhileobject *lz)
1131{
1132 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1133 lz->func, lz->it, lz->start);
1134}
1135
1136static PyObject *
1137dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1138{
1139 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001140 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001141 return NULL;
1142 lz->start = start;
1143 Py_RETURN_NONE;
1144}
1145
1146static PyMethodDef dropwhile_methods[] = {
1147 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1148 reduce_doc},
1149 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1150 setstate_doc},
1151 {NULL, NULL} /* sentinel */
1152};
1153
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154PyDoc_STRVAR(dropwhile_doc,
1155"dropwhile(predicate, iterable) --> dropwhile object\n\
1156\n\
1157Drop items from the iterable while predicate(item) is true.\n\
1158Afterwards, return every element until the iterable is exhausted.");
1159
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001160static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001161 PyVarObject_HEAD_INIT(NULL, 0)
1162 "itertools.dropwhile", /* tp_name */
1163 sizeof(dropwhileobject), /* tp_basicsize */
1164 0, /* tp_itemsize */
1165 /* methods */
1166 (destructor)dropwhile_dealloc, /* tp_dealloc */
1167 0, /* tp_print */
1168 0, /* tp_getattr */
1169 0, /* tp_setattr */
1170 0, /* tp_reserved */
1171 0, /* tp_repr */
1172 0, /* tp_as_number */
1173 0, /* tp_as_sequence */
1174 0, /* tp_as_mapping */
1175 0, /* tp_hash */
1176 0, /* tp_call */
1177 0, /* tp_str */
1178 PyObject_GenericGetAttr, /* tp_getattro */
1179 0, /* tp_setattro */
1180 0, /* tp_as_buffer */
1181 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1182 Py_TPFLAGS_BASETYPE, /* tp_flags */
1183 dropwhile_doc, /* tp_doc */
1184 (traverseproc)dropwhile_traverse, /* tp_traverse */
1185 0, /* tp_clear */
1186 0, /* tp_richcompare */
1187 0, /* tp_weaklistoffset */
1188 PyObject_SelfIter, /* tp_iter */
1189 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001190 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001191 0, /* tp_members */
1192 0, /* tp_getset */
1193 0, /* tp_base */
1194 0, /* tp_dict */
1195 0, /* tp_descr_get */
1196 0, /* tp_descr_set */
1197 0, /* tp_dictoffset */
1198 0, /* tp_init */
1199 0, /* tp_alloc */
1200 dropwhile_new, /* tp_new */
1201 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202};
1203
1204
1205/* takewhile object **********************************************************/
1206
1207typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001208 PyObject_HEAD
1209 PyObject *func;
1210 PyObject *it;
1211 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001212} takewhileobject;
1213
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001214static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215
1216static PyObject *
1217takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1218{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001219 PyObject *func, *seq;
1220 PyObject *it;
1221 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001223 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1224 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001225
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001226 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1227 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001229 /* Get iterator. */
1230 it = PyObject_GetIter(seq);
1231 if (it == NULL)
1232 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 /* create takewhileobject structure */
1235 lz = (takewhileobject *)type->tp_alloc(type, 0);
1236 if (lz == NULL) {
1237 Py_DECREF(it);
1238 return NULL;
1239 }
1240 Py_INCREF(func);
1241 lz->func = func;
1242 lz->it = it;
1243 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001245 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246}
1247
1248static void
1249takewhile_dealloc(takewhileobject *lz)
1250{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001251 PyObject_GC_UnTrack(lz);
1252 Py_XDECREF(lz->func);
1253 Py_XDECREF(lz->it);
1254 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001255}
1256
1257static int
1258takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1259{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001260 Py_VISIT(lz->it);
1261 Py_VISIT(lz->func);
1262 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001263}
1264
1265static PyObject *
1266takewhile_next(takewhileobject *lz)
1267{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001268 PyObject *item, *good;
1269 PyObject *it = lz->it;
1270 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001272 if (lz->stop == 1)
1273 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001274
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001275 item = (*Py_TYPE(it)->tp_iternext)(it);
1276 if (item == NULL)
1277 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001278
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001279 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1280 if (good == NULL) {
1281 Py_DECREF(item);
1282 return NULL;
1283 }
1284 ok = PyObject_IsTrue(good);
1285 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001286 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 return item;
1288 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001289 if (ok == 0)
1290 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001291 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292}
1293
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001294static PyObject *
1295takewhile_reduce(takewhileobject *lz)
1296{
1297 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1298 lz->func, lz->it, lz->stop);
1299}
1300
1301static PyObject *
1302takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1303{
1304 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001305 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001306 return NULL;
1307 lz->stop = stop;
1308 Py_RETURN_NONE;
1309}
1310
1311static PyMethodDef takewhile_reduce_methods[] = {
1312 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1313 reduce_doc},
1314 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1315 setstate_doc},
1316 {NULL, NULL} /* sentinel */
1317};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318PyDoc_STRVAR(takewhile_doc,
1319"takewhile(predicate, iterable) --> takewhile object\n\
1320\n\
1321Return successive entries from an iterable as long as the \n\
1322predicate evaluates to true for each entry.");
1323
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001324static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001325 PyVarObject_HEAD_INIT(NULL, 0)
1326 "itertools.takewhile", /* tp_name */
1327 sizeof(takewhileobject), /* tp_basicsize */
1328 0, /* tp_itemsize */
1329 /* methods */
1330 (destructor)takewhile_dealloc, /* tp_dealloc */
1331 0, /* tp_print */
1332 0, /* tp_getattr */
1333 0, /* tp_setattr */
1334 0, /* tp_reserved */
1335 0, /* tp_repr */
1336 0, /* tp_as_number */
1337 0, /* tp_as_sequence */
1338 0, /* tp_as_mapping */
1339 0, /* tp_hash */
1340 0, /* tp_call */
1341 0, /* tp_str */
1342 PyObject_GenericGetAttr, /* tp_getattro */
1343 0, /* tp_setattro */
1344 0, /* tp_as_buffer */
1345 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1346 Py_TPFLAGS_BASETYPE, /* tp_flags */
1347 takewhile_doc, /* tp_doc */
1348 (traverseproc)takewhile_traverse, /* tp_traverse */
1349 0, /* tp_clear */
1350 0, /* tp_richcompare */
1351 0, /* tp_weaklistoffset */
1352 PyObject_SelfIter, /* tp_iter */
1353 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001354 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001355 0, /* tp_members */
1356 0, /* tp_getset */
1357 0, /* tp_base */
1358 0, /* tp_dict */
1359 0, /* tp_descr_get */
1360 0, /* tp_descr_set */
1361 0, /* tp_dictoffset */
1362 0, /* tp_init */
1363 0, /* tp_alloc */
1364 takewhile_new, /* tp_new */
1365 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366};
1367
1368
1369/* islice object ************************************************************/
1370
1371typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001372 PyObject_HEAD
1373 PyObject *it;
1374 Py_ssize_t next;
1375 Py_ssize_t stop;
1376 Py_ssize_t step;
1377 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378} isliceobject;
1379
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001380static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381
1382static PyObject *
1383islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1384{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001385 PyObject *seq;
1386 Py_ssize_t start=0, stop=-1, step=1;
1387 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1388 Py_ssize_t numargs;
1389 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001391 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1392 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001393
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001394 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1395 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001396
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001397 numargs = PyTuple_Size(args);
1398 if (numargs == 2) {
1399 if (a1 != Py_None) {
1400 stop = PyLong_AsSsize_t(a1);
1401 if (stop == -1) {
1402 if (PyErr_Occurred())
1403 PyErr_Clear();
1404 PyErr_SetString(PyExc_ValueError,
1405 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1406 return NULL;
1407 }
1408 }
1409 } else {
1410 if (a1 != Py_None)
1411 start = PyLong_AsSsize_t(a1);
1412 if (start == -1 && PyErr_Occurred())
1413 PyErr_Clear();
1414 if (a2 != Py_None) {
1415 stop = PyLong_AsSsize_t(a2);
1416 if (stop == -1) {
1417 if (PyErr_Occurred())
1418 PyErr_Clear();
1419 PyErr_SetString(PyExc_ValueError,
1420 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1421 return NULL;
1422 }
1423 }
1424 }
1425 if (start<0 || stop<-1) {
1426 PyErr_SetString(PyExc_ValueError,
1427 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1428 return NULL;
1429 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001431 if (a3 != NULL) {
1432 if (a3 != Py_None)
1433 step = PyLong_AsSsize_t(a3);
1434 if (step == -1 && PyErr_Occurred())
1435 PyErr_Clear();
1436 }
1437 if (step<1) {
1438 PyErr_SetString(PyExc_ValueError,
1439 "Step for islice() must be a positive integer or None.");
1440 return NULL;
1441 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001443 /* Get iterator. */
1444 it = PyObject_GetIter(seq);
1445 if (it == NULL)
1446 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001448 /* create isliceobject structure */
1449 lz = (isliceobject *)type->tp_alloc(type, 0);
1450 if (lz == NULL) {
1451 Py_DECREF(it);
1452 return NULL;
1453 }
1454 lz->it = it;
1455 lz->next = start;
1456 lz->stop = stop;
1457 lz->step = step;
1458 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001460 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461}
1462
1463static void
1464islice_dealloc(isliceobject *lz)
1465{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001466 PyObject_GC_UnTrack(lz);
1467 Py_XDECREF(lz->it);
1468 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469}
1470
1471static int
1472islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1473{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001474 Py_VISIT(lz->it);
1475 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476}
1477
1478static PyObject *
1479islice_next(isliceobject *lz)
1480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001481 PyObject *item;
1482 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001483 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001484 Py_ssize_t oldnext;
1485 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001487 if (it == NULL)
1488 return NULL;
1489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001490 iternext = *Py_TYPE(it)->tp_iternext;
1491 while (lz->cnt < lz->next) {
1492 item = iternext(it);
1493 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001494 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001495 Py_DECREF(item);
1496 lz->cnt++;
1497 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001498 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001499 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001500 item = iternext(it);
1501 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001502 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001503 lz->cnt++;
1504 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001505 /* The (size_t) cast below avoids the danger of undefined
1506 behaviour from signed integer overflow. */
1507 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001508 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1509 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001510 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001511
1512empty:
1513 Py_CLEAR(lz->it);
1514 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001515}
1516
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001517static PyObject *
1518islice_reduce(isliceobject *lz)
1519{
1520 /* When unpickled, generate a new object with the same bounds,
1521 * then 'setstate' with the next and count
1522 */
1523 PyObject *stop;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001524 if (lz->it == NULL) {
1525 PyObject *empty_list;
1526 PyObject *empty_it;
1527 empty_list = PyList_New(0);
1528 if (empty_list == NULL)
1529 return NULL;
1530 empty_it = PyObject_GetIter(empty_list);
1531 Py_DECREF(empty_list);
1532 if (empty_it == NULL)
1533 return NULL;
1534 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1535 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001536 if (lz->stop == -1) {
1537 stop = Py_None;
1538 Py_INCREF(stop);
1539 } else {
1540 stop = PyLong_FromSsize_t(lz->stop);
1541 if (stop == NULL)
1542 return NULL;
1543 }
1544 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1545 lz->it, lz->next, stop, lz->step,
1546 lz->cnt);
1547}
1548
1549static PyObject *
1550islice_setstate(isliceobject *lz, PyObject *state)
1551{
1552 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1553 if (cnt == -1 && PyErr_Occurred())
1554 return NULL;
1555 lz->cnt = cnt;
1556 Py_RETURN_NONE;
1557}
1558
1559static PyMethodDef islice_methods[] = {
1560 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1561 reduce_doc},
1562 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1563 setstate_doc},
1564 {NULL, NULL} /* sentinel */
1565};
1566
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001567PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001568"islice(iterable, stop) --> islice object\n\
1569islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001570\n\
1571Return an iterator whose next() method returns selected values from an\n\
1572iterable. If start is specified, will skip all preceding elements;\n\
1573otherwise, start defaults to zero. Step defaults to one. If\n\
1574specified as another value, step determines how many values are \n\
1575skipped between successive calls. Works like a slice() on a list\n\
1576but returns an iterator.");
1577
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001578static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001579 PyVarObject_HEAD_INIT(NULL, 0)
1580 "itertools.islice", /* tp_name */
1581 sizeof(isliceobject), /* tp_basicsize */
1582 0, /* tp_itemsize */
1583 /* methods */
1584 (destructor)islice_dealloc, /* tp_dealloc */
1585 0, /* tp_print */
1586 0, /* tp_getattr */
1587 0, /* tp_setattr */
1588 0, /* tp_reserved */
1589 0, /* tp_repr */
1590 0, /* tp_as_number */
1591 0, /* tp_as_sequence */
1592 0, /* tp_as_mapping */
1593 0, /* tp_hash */
1594 0, /* tp_call */
1595 0, /* tp_str */
1596 PyObject_GenericGetAttr, /* tp_getattro */
1597 0, /* tp_setattro */
1598 0, /* tp_as_buffer */
1599 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1600 Py_TPFLAGS_BASETYPE, /* tp_flags */
1601 islice_doc, /* tp_doc */
1602 (traverseproc)islice_traverse, /* tp_traverse */
1603 0, /* tp_clear */
1604 0, /* tp_richcompare */
1605 0, /* tp_weaklistoffset */
1606 PyObject_SelfIter, /* tp_iter */
1607 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001608 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001609 0, /* tp_members */
1610 0, /* tp_getset */
1611 0, /* tp_base */
1612 0, /* tp_dict */
1613 0, /* tp_descr_get */
1614 0, /* tp_descr_set */
1615 0, /* tp_dictoffset */
1616 0, /* tp_init */
1617 0, /* tp_alloc */
1618 islice_new, /* tp_new */
1619 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001620};
1621
1622
1623/* starmap object ************************************************************/
1624
1625typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001626 PyObject_HEAD
1627 PyObject *func;
1628 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001629} starmapobject;
1630
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001631static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001632
1633static PyObject *
1634starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1635{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001636 PyObject *func, *seq;
1637 PyObject *it;
1638 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001640 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1641 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001643 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1644 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001646 /* Get iterator. */
1647 it = PyObject_GetIter(seq);
1648 if (it == NULL)
1649 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 /* create starmapobject structure */
1652 lz = (starmapobject *)type->tp_alloc(type, 0);
1653 if (lz == NULL) {
1654 Py_DECREF(it);
1655 return NULL;
1656 }
1657 Py_INCREF(func);
1658 lz->func = func;
1659 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001661 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662}
1663
1664static void
1665starmap_dealloc(starmapobject *lz)
1666{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001667 PyObject_GC_UnTrack(lz);
1668 Py_XDECREF(lz->func);
1669 Py_XDECREF(lz->it);
1670 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671}
1672
1673static int
1674starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1675{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001676 Py_VISIT(lz->it);
1677 Py_VISIT(lz->func);
1678 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679}
1680
1681static PyObject *
1682starmap_next(starmapobject *lz)
1683{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001684 PyObject *args;
1685 PyObject *result;
1686 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001688 args = (*Py_TYPE(it)->tp_iternext)(it);
1689 if (args == NULL)
1690 return NULL;
1691 if (!PyTuple_CheckExact(args)) {
1692 PyObject *newargs = PySequence_Tuple(args);
1693 Py_DECREF(args);
1694 if (newargs == NULL)
1695 return NULL;
1696 args = newargs;
1697 }
1698 result = PyObject_Call(lz->func, args, NULL);
1699 Py_DECREF(args);
1700 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001701}
1702
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001703static PyObject *
1704starmap_reduce(starmapobject *lz)
1705{
1706 /* Just pickle the iterator */
1707 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1708}
1709
1710static PyMethodDef starmap_methods[] = {
1711 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1712 reduce_doc},
1713 {NULL, NULL} /* sentinel */
1714};
1715
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716PyDoc_STRVAR(starmap_doc,
1717"starmap(function, sequence) --> starmap object\n\
1718\n\
1719Return an iterator whose values are returned from the function evaluated\n\
Serhiy Storchakad65c9492015-11-02 14:10:23 +02001720with an argument tuple taken from the given sequence.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001721
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001722static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001723 PyVarObject_HEAD_INIT(NULL, 0)
1724 "itertools.starmap", /* tp_name */
1725 sizeof(starmapobject), /* tp_basicsize */
1726 0, /* tp_itemsize */
1727 /* methods */
1728 (destructor)starmap_dealloc, /* tp_dealloc */
1729 0, /* tp_print */
1730 0, /* tp_getattr */
1731 0, /* tp_setattr */
1732 0, /* tp_reserved */
1733 0, /* tp_repr */
1734 0, /* tp_as_number */
1735 0, /* tp_as_sequence */
1736 0, /* tp_as_mapping */
1737 0, /* tp_hash */
1738 0, /* tp_call */
1739 0, /* tp_str */
1740 PyObject_GenericGetAttr, /* tp_getattro */
1741 0, /* tp_setattro */
1742 0, /* tp_as_buffer */
1743 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1744 Py_TPFLAGS_BASETYPE, /* tp_flags */
1745 starmap_doc, /* tp_doc */
1746 (traverseproc)starmap_traverse, /* tp_traverse */
1747 0, /* tp_clear */
1748 0, /* tp_richcompare */
1749 0, /* tp_weaklistoffset */
1750 PyObject_SelfIter, /* tp_iter */
1751 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001752 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001753 0, /* tp_members */
1754 0, /* tp_getset */
1755 0, /* tp_base */
1756 0, /* tp_dict */
1757 0, /* tp_descr_get */
1758 0, /* tp_descr_set */
1759 0, /* tp_dictoffset */
1760 0, /* tp_init */
1761 0, /* tp_alloc */
1762 starmap_new, /* tp_new */
1763 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764};
1765
1766
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001767/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001768
1769typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001770 PyObject_HEAD
1771 PyObject *source; /* Iterator over input iterables */
1772 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001773} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001774
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001775static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001776
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001777static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001778chain_new_internal(PyTypeObject *type, PyObject *source)
1779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001780 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001781
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001782 lz = (chainobject *)type->tp_alloc(type, 0);
1783 if (lz == NULL) {
1784 Py_DECREF(source);
1785 return NULL;
1786 }
1787
1788 lz->source = source;
1789 lz->active = NULL;
1790 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001791}
1792
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001793static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001794chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001795{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001796 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001797
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001798 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1799 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001800
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001801 source = PyObject_GetIter(args);
1802 if (source == NULL)
1803 return NULL;
1804
1805 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001806}
1807
1808static PyObject *
1809chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1810{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001811 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001813 source = PyObject_GetIter(arg);
1814 if (source == NULL)
1815 return NULL;
1816
1817 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001818}
1819
1820static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001821chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001822{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001823 PyObject_GC_UnTrack(lz);
1824 Py_XDECREF(lz->active);
1825 Py_XDECREF(lz->source);
1826 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001827}
1828
Raymond Hettinger2012f172003-02-07 05:32:58 +00001829static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001830chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001831{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001832 Py_VISIT(lz->source);
1833 Py_VISIT(lz->active);
1834 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001835}
1836
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001837static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001838chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001840 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001841
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001842 if (lz->source == NULL)
1843 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001844
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001845 if (lz->active == NULL) {
1846 PyObject *iterable = PyIter_Next(lz->source);
1847 if (iterable == NULL) {
1848 Py_CLEAR(lz->source);
1849 return NULL; /* no more input sources */
1850 }
1851 lz->active = PyObject_GetIter(iterable);
1852 Py_DECREF(iterable);
1853 if (lz->active == NULL) {
1854 Py_CLEAR(lz->source);
1855 return NULL; /* input not iterable */
1856 }
1857 }
1858 item = PyIter_Next(lz->active);
1859 if (item != NULL)
1860 return item;
1861 if (PyErr_Occurred()) {
1862 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1863 PyErr_Clear();
1864 else
1865 return NULL; /* input raised an exception */
1866 }
1867 Py_CLEAR(lz->active);
1868 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001869}
1870
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001871static PyObject *
1872chain_reduce(chainobject *lz)
1873{
1874 if (lz->source) {
1875 /* we can't pickle function objects (itertools.from_iterable) so
1876 * we must use setstate to replace the iterable. One day we
1877 * will fix pickling of functions
1878 */
1879 if (lz->active) {
1880 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1881 } else {
1882 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1883 }
1884 } else {
1885 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1886 }
1887 return NULL;
1888}
1889
1890static PyObject *
1891chain_setstate(chainobject *lz, PyObject *state)
1892{
1893 PyObject *source, *active=NULL;
1894 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1895 return NULL;
1896
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02001897 Py_INCREF(source);
1898 Py_SETREF(lz->source, source);
1899 Py_XINCREF(active);
1900 Py_SETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001901 Py_RETURN_NONE;
1902}
1903
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001904PyDoc_STRVAR(chain_doc,
1905"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001906\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001907Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001908first iterable until it is exhausted, then elements from the next\n\
1909iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001910
Christian Heimesf16baeb2008-02-29 14:57:44 +00001911PyDoc_STRVAR(chain_from_iterable_doc,
1912"chain.from_iterable(iterable) --> chain object\n\
1913\n\
1914Alternate chain() contructor taking a single iterable argument\n\
1915that evaluates lazily.");
1916
1917static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001918 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1919 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001920 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1921 reduce_doc},
1922 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1923 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001924 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001925};
1926
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001927static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001928 PyVarObject_HEAD_INIT(NULL, 0)
1929 "itertools.chain", /* tp_name */
1930 sizeof(chainobject), /* tp_basicsize */
1931 0, /* tp_itemsize */
1932 /* methods */
1933 (destructor)chain_dealloc, /* tp_dealloc */
1934 0, /* tp_print */
1935 0, /* tp_getattr */
1936 0, /* tp_setattr */
1937 0, /* tp_reserved */
1938 0, /* tp_repr */
1939 0, /* tp_as_number */
1940 0, /* tp_as_sequence */
1941 0, /* tp_as_mapping */
1942 0, /* tp_hash */
1943 0, /* tp_call */
1944 0, /* tp_str */
1945 PyObject_GenericGetAttr, /* tp_getattro */
1946 0, /* tp_setattro */
1947 0, /* tp_as_buffer */
1948 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1949 Py_TPFLAGS_BASETYPE, /* tp_flags */
1950 chain_doc, /* tp_doc */
1951 (traverseproc)chain_traverse, /* tp_traverse */
1952 0, /* tp_clear */
1953 0, /* tp_richcompare */
1954 0, /* tp_weaklistoffset */
1955 PyObject_SelfIter, /* tp_iter */
1956 (iternextfunc)chain_next, /* tp_iternext */
1957 chain_methods, /* tp_methods */
1958 0, /* tp_members */
1959 0, /* tp_getset */
1960 0, /* tp_base */
1961 0, /* tp_dict */
1962 0, /* tp_descr_get */
1963 0, /* tp_descr_set */
1964 0, /* tp_dictoffset */
1965 0, /* tp_init */
1966 0, /* tp_alloc */
1967 chain_new, /* tp_new */
1968 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001969};
1970
1971
Christian Heimesc3f30c42008-02-22 16:37:40 +00001972/* product object ************************************************************/
1973
1974typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001975 PyObject_HEAD
1976 PyObject *pools; /* tuple of pool tuples */
1977 Py_ssize_t *indices; /* one index per pool */
1978 PyObject *result; /* most recently returned result tuple */
1979 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001980} productobject;
1981
1982static PyTypeObject product_type;
1983
1984static PyObject *
1985product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1986{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001987 productobject *lz;
1988 Py_ssize_t nargs, npools, repeat=1;
1989 PyObject *pools = NULL;
1990 Py_ssize_t *indices = NULL;
1991 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001992
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001993 if (kwds != NULL) {
1994 char *kwlist[] = {"repeat", 0};
1995 PyObject *tmpargs = PyTuple_New(0);
1996 if (tmpargs == NULL)
1997 return NULL;
1998 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1999 Py_DECREF(tmpargs);
2000 return NULL;
2001 }
2002 Py_DECREF(tmpargs);
2003 if (repeat < 0) {
2004 PyErr_SetString(PyExc_ValueError,
2005 "repeat argument cannot be negative");
2006 return NULL;
2007 }
2008 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002009
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002010 assert(PyTuple_CheckExact(args));
2011 if (repeat == 0) {
2012 nargs = 0;
2013 } else {
2014 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002015 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002016 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2017 return NULL;
2018 }
2019 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002020 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002021
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002022 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002023 if (indices == NULL) {
2024 PyErr_NoMemory();
2025 goto error;
2026 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002027
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002028 pools = PyTuple_New(npools);
2029 if (pools == NULL)
2030 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002032 for (i=0; i < nargs ; ++i) {
2033 PyObject *item = PyTuple_GET_ITEM(args, i);
2034 PyObject *pool = PySequence_Tuple(item);
2035 if (pool == NULL)
2036 goto error;
2037 PyTuple_SET_ITEM(pools, i, pool);
2038 indices[i] = 0;
2039 }
2040 for ( ; i < npools; ++i) {
2041 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2042 Py_INCREF(pool);
2043 PyTuple_SET_ITEM(pools, i, pool);
2044 indices[i] = 0;
2045 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002046
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002047 /* create productobject structure */
2048 lz = (productobject *)type->tp_alloc(type, 0);
2049 if (lz == NULL)
2050 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002051
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002052 lz->pools = pools;
2053 lz->indices = indices;
2054 lz->result = NULL;
2055 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002057 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002058
2059error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002060 if (indices != NULL)
2061 PyMem_Free(indices);
2062 Py_XDECREF(pools);
2063 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002064}
2065
2066static void
2067product_dealloc(productobject *lz)
2068{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002069 PyObject_GC_UnTrack(lz);
2070 Py_XDECREF(lz->pools);
2071 Py_XDECREF(lz->result);
2072 if (lz->indices != NULL)
2073 PyMem_Free(lz->indices);
2074 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002075}
2076
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002077static PyObject *
2078product_sizeof(productobject *lz, void *unused)
2079{
2080 Py_ssize_t res;
2081
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002082 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002083 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2084 return PyLong_FromSsize_t(res);
2085}
2086
2087PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2088
Christian Heimesc3f30c42008-02-22 16:37:40 +00002089static int
2090product_traverse(productobject *lz, visitproc visit, void *arg)
2091{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002092 Py_VISIT(lz->pools);
2093 Py_VISIT(lz->result);
2094 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002095}
2096
2097static PyObject *
2098product_next(productobject *lz)
2099{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002100 PyObject *pool;
2101 PyObject *elem;
2102 PyObject *oldelem;
2103 PyObject *pools = lz->pools;
2104 PyObject *result = lz->result;
2105 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2106 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002107
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002108 if (lz->stopped)
2109 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002110
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002111 if (result == NULL) {
2112 /* On the first pass, return an initial tuple filled with the
2113 first element from each pool. */
2114 result = PyTuple_New(npools);
2115 if (result == NULL)
2116 goto empty;
2117 lz->result = result;
2118 for (i=0; i < npools; i++) {
2119 pool = PyTuple_GET_ITEM(pools, i);
2120 if (PyTuple_GET_SIZE(pool) == 0)
2121 goto empty;
2122 elem = PyTuple_GET_ITEM(pool, 0);
2123 Py_INCREF(elem);
2124 PyTuple_SET_ITEM(result, i, elem);
2125 }
2126 } else {
2127 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002129 /* Copy the previous result tuple or re-use it if available */
2130 if (Py_REFCNT(result) > 1) {
2131 PyObject *old_result = result;
2132 result = PyTuple_New(npools);
2133 if (result == NULL)
2134 goto empty;
2135 lz->result = result;
2136 for (i=0; i < npools; i++) {
2137 elem = PyTuple_GET_ITEM(old_result, i);
2138 Py_INCREF(elem);
2139 PyTuple_SET_ITEM(result, i, elem);
2140 }
2141 Py_DECREF(old_result);
2142 }
2143 /* Now, we've got the only copy so we can update it in-place */
2144 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002146 /* Update the pool indices right-to-left. Only advance to the
2147 next pool when the previous one rolls-over */
2148 for (i=npools-1 ; i >= 0 ; i--) {
2149 pool = PyTuple_GET_ITEM(pools, i);
2150 indices[i]++;
2151 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2152 /* Roll-over and advance to next pool */
2153 indices[i] = 0;
2154 elem = PyTuple_GET_ITEM(pool, 0);
2155 Py_INCREF(elem);
2156 oldelem = PyTuple_GET_ITEM(result, i);
2157 PyTuple_SET_ITEM(result, i, elem);
2158 Py_DECREF(oldelem);
2159 } else {
2160 /* No rollover. Just increment and stop here. */
2161 elem = PyTuple_GET_ITEM(pool, indices[i]);
2162 Py_INCREF(elem);
2163 oldelem = PyTuple_GET_ITEM(result, i);
2164 PyTuple_SET_ITEM(result, i, elem);
2165 Py_DECREF(oldelem);
2166 break;
2167 }
2168 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002169
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002170 /* If i is negative, then the indices have all rolled-over
2171 and we're done. */
2172 if (i < 0)
2173 goto empty;
2174 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002175
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002176 Py_INCREF(result);
2177 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002178
2179empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 lz->stopped = 1;
2181 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002182}
2183
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002184static PyObject *
2185product_reduce(productobject *lz)
2186{
2187 if (lz->stopped) {
2188 return Py_BuildValue("O(())", Py_TYPE(lz));
2189 } else if (lz->result == NULL) {
2190 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2191 } else {
2192 PyObject *indices;
2193 Py_ssize_t n, i;
2194
2195 /* we must pickle the indices use them for setstate, and
2196 * additionally indicate that the iterator has started
2197 */
2198 n = PyTuple_GET_SIZE(lz->pools);
2199 indices = PyTuple_New(n);
2200 if (indices == NULL)
2201 return NULL;
2202 for (i=0; i<n; i++){
2203 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2204 if (!index) {
2205 Py_DECREF(indices);
2206 return NULL;
2207 }
2208 PyTuple_SET_ITEM(indices, i, index);
2209 }
2210 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2211 }
2212}
2213
2214static PyObject *
2215product_setstate(productobject *lz, PyObject *state)
2216{
2217 PyObject *result;
2218 Py_ssize_t n, i;
2219
2220 n = PyTuple_GET_SIZE(lz->pools);
2221 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2222 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2223 return NULL;
2224 }
2225 for (i=0; i<n; i++)
2226 {
2227 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2228 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002229 PyObject* pool;
2230 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002231 if (index < 0 && PyErr_Occurred())
2232 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002233 pool = PyTuple_GET_ITEM(lz->pools, i);
2234 poolsize = PyTuple_GET_SIZE(pool);
2235 if (poolsize == 0) {
2236 lz->stopped = 1;
2237 Py_RETURN_NONE;
2238 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002239 /* clamp the index */
2240 if (index < 0)
2241 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002242 else if (index > poolsize-1)
2243 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002244 lz->indices[i] = index;
2245 }
2246
2247 result = PyTuple_New(n);
2248 if (!result)
2249 return NULL;
2250 for (i=0; i<n; i++) {
2251 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2252 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2253 Py_INCREF(element);
2254 PyTuple_SET_ITEM(result, i, element);
2255 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002256 Py_SETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002257 Py_RETURN_NONE;
2258}
2259
2260static PyMethodDef product_methods[] = {
2261 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2262 reduce_doc},
2263 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2264 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002265 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2266 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002267 {NULL, NULL} /* sentinel */
2268};
2269
Christian Heimesc3f30c42008-02-22 16:37:40 +00002270PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002271"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002272\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002273Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002274For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2275The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2276cycle in a manner similar to an odometer (with the rightmost element changing\n\
2277on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002278To compute the product of an iterable with itself, specify the number\n\
2279of repetitions with the optional repeat keyword argument. For example,\n\
2280product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002281product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2282product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2283
2284static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002285 PyVarObject_HEAD_INIT(NULL, 0)
2286 "itertools.product", /* tp_name */
2287 sizeof(productobject), /* tp_basicsize */
2288 0, /* tp_itemsize */
2289 /* methods */
2290 (destructor)product_dealloc, /* tp_dealloc */
2291 0, /* tp_print */
2292 0, /* tp_getattr */
2293 0, /* tp_setattr */
2294 0, /* tp_reserved */
2295 0, /* tp_repr */
2296 0, /* tp_as_number */
2297 0, /* tp_as_sequence */
2298 0, /* tp_as_mapping */
2299 0, /* tp_hash */
2300 0, /* tp_call */
2301 0, /* tp_str */
2302 PyObject_GenericGetAttr, /* tp_getattro */
2303 0, /* tp_setattro */
2304 0, /* tp_as_buffer */
2305 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2306 Py_TPFLAGS_BASETYPE, /* tp_flags */
2307 product_doc, /* tp_doc */
2308 (traverseproc)product_traverse, /* tp_traverse */
2309 0, /* tp_clear */
2310 0, /* tp_richcompare */
2311 0, /* tp_weaklistoffset */
2312 PyObject_SelfIter, /* tp_iter */
2313 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002314 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002315 0, /* tp_members */
2316 0, /* tp_getset */
2317 0, /* tp_base */
2318 0, /* tp_dict */
2319 0, /* tp_descr_get */
2320 0, /* tp_descr_set */
2321 0, /* tp_dictoffset */
2322 0, /* tp_init */
2323 0, /* tp_alloc */
2324 product_new, /* tp_new */
2325 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002326};
2327
2328
Christian Heimes380f7f22008-02-28 11:19:05 +00002329/* combinations object ************************************************************/
2330
2331typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002332 PyObject_HEAD
2333 PyObject *pool; /* input converted to a tuple */
2334 Py_ssize_t *indices; /* one index per result element */
2335 PyObject *result; /* most recently returned result tuple */
2336 Py_ssize_t r; /* size of result tuple */
2337 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002338} combinationsobject;
2339
2340static PyTypeObject combinations_type;
2341
2342static PyObject *
2343combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2344{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002345 combinationsobject *co;
2346 Py_ssize_t n;
2347 Py_ssize_t r;
2348 PyObject *pool = NULL;
2349 PyObject *iterable = NULL;
2350 Py_ssize_t *indices = NULL;
2351 Py_ssize_t i;
2352 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002353
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002354 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2355 &iterable, &r))
2356 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002357
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002358 pool = PySequence_Tuple(iterable);
2359 if (pool == NULL)
2360 goto error;
2361 n = PyTuple_GET_SIZE(pool);
2362 if (r < 0) {
2363 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2364 goto error;
2365 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002366
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002367 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002368 if (indices == NULL) {
2369 PyErr_NoMemory();
2370 goto error;
2371 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002372
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002373 for (i=0 ; i<r ; i++)
2374 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002375
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 /* create combinationsobject structure */
2377 co = (combinationsobject *)type->tp_alloc(type, 0);
2378 if (co == NULL)
2379 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002380
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002381 co->pool = pool;
2382 co->indices = indices;
2383 co->result = NULL;
2384 co->r = r;
2385 co->stopped = r > n ? 1 : 0;
2386
2387 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002388
2389error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002390 if (indices != NULL)
2391 PyMem_Free(indices);
2392 Py_XDECREF(pool);
2393 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002394}
2395
2396static void
2397combinations_dealloc(combinationsobject *co)
2398{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002399 PyObject_GC_UnTrack(co);
2400 Py_XDECREF(co->pool);
2401 Py_XDECREF(co->result);
2402 if (co->indices != NULL)
2403 PyMem_Free(co->indices);
2404 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002405}
2406
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002407static PyObject *
2408combinations_sizeof(combinationsobject *co, void *unused)
2409{
2410 Py_ssize_t res;
2411
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002412 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002413 res += co->r * sizeof(Py_ssize_t);
2414 return PyLong_FromSsize_t(res);
2415}
2416
Christian Heimes380f7f22008-02-28 11:19:05 +00002417static int
2418combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2419{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002420 Py_VISIT(co->pool);
2421 Py_VISIT(co->result);
2422 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002423}
2424
2425static PyObject *
2426combinations_next(combinationsobject *co)
2427{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002428 PyObject *elem;
2429 PyObject *oldelem;
2430 PyObject *pool = co->pool;
2431 Py_ssize_t *indices = co->indices;
2432 PyObject *result = co->result;
2433 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2434 Py_ssize_t r = co->r;
2435 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002437 if (co->stopped)
2438 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002439
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002440 if (result == NULL) {
2441 /* On the first pass, initialize result tuple using the indices */
2442 result = PyTuple_New(r);
2443 if (result == NULL)
2444 goto empty;
2445 co->result = result;
2446 for (i=0; i<r ; i++) {
2447 index = indices[i];
2448 elem = PyTuple_GET_ITEM(pool, index);
2449 Py_INCREF(elem);
2450 PyTuple_SET_ITEM(result, i, elem);
2451 }
2452 } else {
2453 /* Copy the previous result tuple or re-use it if available */
2454 if (Py_REFCNT(result) > 1) {
2455 PyObject *old_result = result;
2456 result = PyTuple_New(r);
2457 if (result == NULL)
2458 goto empty;
2459 co->result = result;
2460 for (i=0; i<r ; i++) {
2461 elem = PyTuple_GET_ITEM(old_result, i);
2462 Py_INCREF(elem);
2463 PyTuple_SET_ITEM(result, i, elem);
2464 }
2465 Py_DECREF(old_result);
2466 }
2467 /* Now, we've got the only copy so we can update it in-place
2468 * CPython's empty tuple is a singleton and cached in
2469 * PyTuple's freelist.
2470 */
2471 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002472
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002473 /* Scan indices right-to-left until finding one that is not
2474 at its maximum (i + n - r). */
2475 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2476 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002477
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002478 /* If i is negative, then the indices are all at
2479 their maximum value and we're done. */
2480 if (i < 0)
2481 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002482
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002483 /* Increment the current index which we know is not at its
2484 maximum. Then move back to the right setting each index
2485 to its lowest possible value (one higher than the index
2486 to its left -- this maintains the sort order invariant). */
2487 indices[i]++;
2488 for (j=i+1 ; j<r ; j++)
2489 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002490
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002491 /* Update the result tuple for the new indices
2492 starting with i, the leftmost index that changed */
2493 for ( ; i<r ; i++) {
2494 index = indices[i];
2495 elem = PyTuple_GET_ITEM(pool, index);
2496 Py_INCREF(elem);
2497 oldelem = PyTuple_GET_ITEM(result, i);
2498 PyTuple_SET_ITEM(result, i, elem);
2499 Py_DECREF(oldelem);
2500 }
2501 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002502
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002503 Py_INCREF(result);
2504 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002505
2506empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002507 co->stopped = 1;
2508 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002509}
2510
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002511static PyObject *
2512combinations_reduce(combinationsobject *lz)
2513{
2514 if (lz->result == NULL) {
2515 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2516 } else if (lz->stopped) {
2517 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2518 } else {
2519 PyObject *indices;
2520 Py_ssize_t i;
2521
2522 /* we must pickle the indices and use them for setstate */
2523 indices = PyTuple_New(lz->r);
2524 if (!indices)
2525 return NULL;
2526 for (i=0; i<lz->r; i++)
2527 {
2528 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2529 if (!index) {
2530 Py_DECREF(indices);
2531 return NULL;
2532 }
2533 PyTuple_SET_ITEM(indices, i, index);
2534 }
2535
2536 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2537 }
2538}
2539
2540static PyObject *
2541combinations_setstate(combinationsobject *lz, PyObject *state)
2542{
2543 PyObject *result;
2544 Py_ssize_t i;
2545 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2546
2547 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2548 {
2549 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2550 return NULL;
2551 }
2552
2553 for (i=0; i<lz->r; i++)
2554 {
2555 Py_ssize_t max;
2556 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2557 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2558 if (index == -1 && PyErr_Occurred())
2559 return NULL; /* not an integer */
2560 max = i + n - lz->r;
2561 /* clamp the index (beware of negative max) */
2562 if (index > max)
2563 index = max;
2564 if (index < 0)
2565 index = 0;
2566 lz->indices[i] = index;
2567 }
2568
2569 result = PyTuple_New(lz->r);
2570 if (result == NULL)
2571 return NULL;
2572 for (i=0; i<lz->r; i++) {
2573 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2574 Py_INCREF(element);
2575 PyTuple_SET_ITEM(result, i, element);
2576 }
2577
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002578 Py_SETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002579 Py_RETURN_NONE;
2580}
2581
2582static PyMethodDef combinations_methods[] = {
2583 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2584 reduce_doc},
2585 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2586 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002587 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2588 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002589 {NULL, NULL} /* sentinel */
2590};
2591
Christian Heimes380f7f22008-02-28 11:19:05 +00002592PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002593"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002594\n\
2595Return successive r-length combinations of elements in the iterable.\n\n\
2596combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2597
2598static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002599 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002600 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002601 sizeof(combinationsobject), /* tp_basicsize */
2602 0, /* tp_itemsize */
2603 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002604 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002605 0, /* tp_print */
2606 0, /* tp_getattr */
2607 0, /* tp_setattr */
2608 0, /* tp_reserved */
2609 0, /* tp_repr */
2610 0, /* tp_as_number */
2611 0, /* tp_as_sequence */
2612 0, /* tp_as_mapping */
2613 0, /* tp_hash */
2614 0, /* tp_call */
2615 0, /* tp_str */
2616 PyObject_GenericGetAttr, /* tp_getattro */
2617 0, /* tp_setattro */
2618 0, /* tp_as_buffer */
2619 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2620 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002621 combinations_doc, /* tp_doc */
2622 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002623 0, /* tp_clear */
2624 0, /* tp_richcompare */
2625 0, /* tp_weaklistoffset */
2626 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002627 (iternextfunc)combinations_next, /* tp_iternext */
2628 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002629 0, /* tp_members */
2630 0, /* tp_getset */
2631 0, /* tp_base */
2632 0, /* tp_dict */
2633 0, /* tp_descr_get */
2634 0, /* tp_descr_set */
2635 0, /* tp_dictoffset */
2636 0, /* tp_init */
2637 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002638 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002639 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002640};
2641
2642
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002643/* combinations with replacement object *******************************************/
2644
2645/* Equivalent to:
2646
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002647 def combinations_with_replacement(iterable, r):
2648 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2649 # number items returned: (n+r-1)! / r! / (n-1)!
2650 pool = tuple(iterable)
2651 n = len(pool)
2652 indices = [0] * r
2653 yield tuple(pool[i] for i in indices)
2654 while 1:
2655 for i in reversed(range(r)):
2656 if indices[i] != n - 1:
2657 break
2658 else:
2659 return
2660 indices[i:] = [indices[i] + 1] * (r - i)
2661 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002662
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002663 def combinations_with_replacement2(iterable, r):
2664 'Alternate version that filters from product()'
2665 pool = tuple(iterable)
2666 n = len(pool)
2667 for indices in product(range(n), repeat=r):
2668 if sorted(indices) == list(indices):
2669 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002670*/
2671typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002672 PyObject_HEAD
2673 PyObject *pool; /* input converted to a tuple */
2674 Py_ssize_t *indices; /* one index per result element */
2675 PyObject *result; /* most recently returned result tuple */
2676 Py_ssize_t r; /* size of result tuple */
2677 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002678} cwrobject;
2679
2680static PyTypeObject cwr_type;
2681
2682static PyObject *
2683cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002685 cwrobject *co;
2686 Py_ssize_t n;
2687 Py_ssize_t r;
2688 PyObject *pool = NULL;
2689 PyObject *iterable = NULL;
2690 Py_ssize_t *indices = NULL;
2691 Py_ssize_t i;
2692 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002693
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002694 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2695 &iterable, &r))
2696 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 pool = PySequence_Tuple(iterable);
2699 if (pool == NULL)
2700 goto error;
2701 n = PyTuple_GET_SIZE(pool);
2702 if (r < 0) {
2703 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2704 goto error;
2705 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002706
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002707 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 if (indices == NULL) {
2709 PyErr_NoMemory();
2710 goto error;
2711 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002712
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002713 for (i=0 ; i<r ; i++)
2714 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002715
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002716 /* create cwrobject structure */
2717 co = (cwrobject *)type->tp_alloc(type, 0);
2718 if (co == NULL)
2719 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002720
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002721 co->pool = pool;
2722 co->indices = indices;
2723 co->result = NULL;
2724 co->r = r;
2725 co->stopped = !n && r;
2726
2727 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002728
2729error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002730 if (indices != NULL)
2731 PyMem_Free(indices);
2732 Py_XDECREF(pool);
2733 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002734}
2735
2736static void
2737cwr_dealloc(cwrobject *co)
2738{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002739 PyObject_GC_UnTrack(co);
2740 Py_XDECREF(co->pool);
2741 Py_XDECREF(co->result);
2742 if (co->indices != NULL)
2743 PyMem_Free(co->indices);
2744 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002745}
2746
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002747static PyObject *
2748cwr_sizeof(cwrobject *co, void *unused)
2749{
2750 Py_ssize_t res;
2751
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002752 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002753 res += co->r * sizeof(Py_ssize_t);
2754 return PyLong_FromSsize_t(res);
2755}
2756
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002757static int
2758cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2759{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 Py_VISIT(co->pool);
2761 Py_VISIT(co->result);
2762 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002763}
2764
2765static PyObject *
2766cwr_next(cwrobject *co)
2767{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002768 PyObject *elem;
2769 PyObject *oldelem;
2770 PyObject *pool = co->pool;
2771 Py_ssize_t *indices = co->indices;
2772 PyObject *result = co->result;
2773 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2774 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002775 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002776
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002777 if (co->stopped)
2778 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002779
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002780 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002781 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002782 result = PyTuple_New(r);
2783 if (result == NULL)
2784 goto empty;
2785 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002786 if (n > 0) {
2787 elem = PyTuple_GET_ITEM(pool, 0);
2788 for (i=0; i<r ; i++) {
2789 assert(indices[i] == 0);
2790 Py_INCREF(elem);
2791 PyTuple_SET_ITEM(result, i, elem);
2792 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 }
2794 } else {
2795 /* Copy the previous result tuple or re-use it if available */
2796 if (Py_REFCNT(result) > 1) {
2797 PyObject *old_result = result;
2798 result = PyTuple_New(r);
2799 if (result == NULL)
2800 goto empty;
2801 co->result = result;
2802 for (i=0; i<r ; i++) {
2803 elem = PyTuple_GET_ITEM(old_result, i);
2804 Py_INCREF(elem);
2805 PyTuple_SET_ITEM(result, i, elem);
2806 }
2807 Py_DECREF(old_result);
2808 }
2809 /* Now, we've got the only copy so we can update it in-place CPython's
2810 empty tuple is a singleton and cached in PyTuple's freelist. */
2811 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002812
Tim Peters9edb1682013-09-03 11:49:31 -05002813 /* Scan indices right-to-left until finding one that is not
2814 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002815 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2816 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002817
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002818 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002819 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002820 if (i < 0)
2821 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002822
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002823 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002824 maximum. Then set all to the right to the same value. */
2825 index = indices[i] + 1;
2826 assert(index < n);
2827 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002828 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002829 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 Py_INCREF(elem);
2831 oldelem = PyTuple_GET_ITEM(result, i);
2832 PyTuple_SET_ITEM(result, i, elem);
2833 Py_DECREF(oldelem);
2834 }
2835 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002836
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002837 Py_INCREF(result);
2838 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002839
2840empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002841 co->stopped = 1;
2842 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002843}
2844
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002845static PyObject *
2846cwr_reduce(cwrobject *lz)
2847{
2848 if (lz->result == NULL) {
2849 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2850 } else if (lz->stopped) {
2851 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2852 } else {
2853 PyObject *indices;
2854 Py_ssize_t i;
2855
2856 /* we must pickle the indices and use them for setstate */
2857 indices = PyTuple_New(lz->r);
2858 if (!indices)
2859 return NULL;
2860 for (i=0; i<lz->r; i++)
2861 {
2862 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2863 if (!index) {
2864 Py_DECREF(indices);
2865 return NULL;
2866 }
2867 PyTuple_SET_ITEM(indices, i, index);
2868 }
2869
2870 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2871 }
2872}
2873
2874static PyObject *
2875cwr_setstate(cwrobject *lz, PyObject *state)
2876{
2877 PyObject *result;
2878 Py_ssize_t n, i;
2879
2880 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2881 {
2882 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2883 return NULL;
2884 }
2885
2886 n = PyTuple_GET_SIZE(lz->pool);
2887 for (i=0; i<lz->r; i++)
2888 {
2889 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2890 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2891 if (index < 0 && PyErr_Occurred())
2892 return NULL; /* not an integer */
2893 /* clamp the index */
2894 if (index < 0)
2895 index = 0;
2896 else if (index > n-1)
2897 index = n-1;
2898 lz->indices[i] = index;
2899 }
2900 result = PyTuple_New(lz->r);
2901 if (result == NULL)
2902 return NULL;
2903 for (i=0; i<lz->r; i++) {
2904 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2905 Py_INCREF(element);
2906 PyTuple_SET_ITEM(result, i, element);
2907 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002908 Py_SETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002909 Py_RETURN_NONE;
2910}
2911
2912static PyMethodDef cwr_methods[] = {
2913 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2914 reduce_doc},
2915 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2916 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002917 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2918 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002919 {NULL, NULL} /* sentinel */
2920};
2921
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002922PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002923"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002924\n\
2925Return successive r-length combinations of elements in the iterable\n\
2926allowing individual elements to have successive repeats.\n\
2927combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2928
2929static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002930 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002931 "itertools.combinations_with_replacement", /* tp_name */
2932 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 0, /* tp_itemsize */
2934 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002935 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002936 0, /* tp_print */
2937 0, /* tp_getattr */
2938 0, /* tp_setattr */
2939 0, /* tp_reserved */
2940 0, /* tp_repr */
2941 0, /* tp_as_number */
2942 0, /* tp_as_sequence */
2943 0, /* tp_as_mapping */
2944 0, /* tp_hash */
2945 0, /* tp_call */
2946 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002947 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002948 0, /* tp_setattro */
2949 0, /* tp_as_buffer */
2950 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002951 Py_TPFLAGS_BASETYPE, /* tp_flags */
2952 cwr_doc, /* tp_doc */
2953 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002954 0, /* tp_clear */
2955 0, /* tp_richcompare */
2956 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002957 PyObject_SelfIter, /* tp_iter */
2958 (iternextfunc)cwr_next, /* tp_iternext */
2959 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002960 0, /* tp_members */
2961 0, /* tp_getset */
2962 0, /* tp_base */
2963 0, /* tp_dict */
2964 0, /* tp_descr_get */
2965 0, /* tp_descr_set */
2966 0, /* tp_dictoffset */
2967 0, /* tp_init */
2968 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002969 cwr_new, /* tp_new */
2970 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002971};
2972
2973
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002974/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002975
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002976def permutations(iterable, r=None):
2977 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2978 pool = tuple(iterable)
2979 n = len(pool)
2980 r = n if r is None else r
2981 indices = range(n)
2982 cycles = range(n-r+1, n+1)[::-1]
2983 yield tuple(pool[i] for i in indices[:r])
2984 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03002985 for i in reversed(range(r)):
2986 cycles[i] -= 1
2987 if cycles[i] == 0:
2988 indices[i:] = indices[i+1:] + indices[i:i+1]
2989 cycles[i] = n - i
2990 else:
2991 j = cycles[i]
2992 indices[i], indices[-j] = indices[-j], indices[i]
2993 yield tuple(pool[i] for i in indices[:r])
2994 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002995 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03002996 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002997*/
2998
2999typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003000 PyObject_HEAD
3001 PyObject *pool; /* input converted to a tuple */
3002 Py_ssize_t *indices; /* one index per element in the pool */
3003 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3004 PyObject *result; /* most recently returned result tuple */
3005 Py_ssize_t r; /* size of result tuple */
3006 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003007} permutationsobject;
3008
3009static PyTypeObject permutations_type;
3010
3011static PyObject *
3012permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3013{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003014 permutationsobject *po;
3015 Py_ssize_t n;
3016 Py_ssize_t r;
3017 PyObject *robj = Py_None;
3018 PyObject *pool = NULL;
3019 PyObject *iterable = NULL;
3020 Py_ssize_t *indices = NULL;
3021 Py_ssize_t *cycles = NULL;
3022 Py_ssize_t i;
3023 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003024
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3026 &iterable, &robj))
3027 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003028
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003029 pool = PySequence_Tuple(iterable);
3030 if (pool == NULL)
3031 goto error;
3032 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003034 r = n;
3035 if (robj != Py_None) {
3036 if (!PyLong_Check(robj)) {
3037 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3038 goto error;
3039 }
3040 r = PyLong_AsSsize_t(robj);
3041 if (r == -1 && PyErr_Occurred())
3042 goto error;
3043 }
3044 if (r < 0) {
3045 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3046 goto error;
3047 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003048
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003049 indices = PyMem_New(Py_ssize_t, n);
3050 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003051 if (indices == NULL || cycles == NULL) {
3052 PyErr_NoMemory();
3053 goto error;
3054 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003055
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003056 for (i=0 ; i<n ; i++)
3057 indices[i] = i;
3058 for (i=0 ; i<r ; i++)
3059 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003060
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 /* create permutationsobject structure */
3062 po = (permutationsobject *)type->tp_alloc(type, 0);
3063 if (po == NULL)
3064 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003065
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 po->pool = pool;
3067 po->indices = indices;
3068 po->cycles = cycles;
3069 po->result = NULL;
3070 po->r = r;
3071 po->stopped = r > n ? 1 : 0;
3072
3073 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003074
3075error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003076 if (indices != NULL)
3077 PyMem_Free(indices);
3078 if (cycles != NULL)
3079 PyMem_Free(cycles);
3080 Py_XDECREF(pool);
3081 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003082}
3083
3084static void
3085permutations_dealloc(permutationsobject *po)
3086{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003087 PyObject_GC_UnTrack(po);
3088 Py_XDECREF(po->pool);
3089 Py_XDECREF(po->result);
3090 PyMem_Free(po->indices);
3091 PyMem_Free(po->cycles);
3092 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003093}
3094
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003095static PyObject *
3096permutations_sizeof(permutationsobject *po, void *unused)
3097{
3098 Py_ssize_t res;
3099
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003100 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003101 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3102 res += po->r * sizeof(Py_ssize_t);
3103 return PyLong_FromSsize_t(res);
3104}
3105
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003106static int
3107permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3108{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003109 Py_VISIT(po->pool);
3110 Py_VISIT(po->result);
3111 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003112}
3113
3114static PyObject *
3115permutations_next(permutationsobject *po)
3116{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003117 PyObject *elem;
3118 PyObject *oldelem;
3119 PyObject *pool = po->pool;
3120 Py_ssize_t *indices = po->indices;
3121 Py_ssize_t *cycles = po->cycles;
3122 PyObject *result = po->result;
3123 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3124 Py_ssize_t r = po->r;
3125 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003127 if (po->stopped)
3128 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003129
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003130 if (result == NULL) {
3131 /* On the first pass, initialize result tuple using the indices */
3132 result = PyTuple_New(r);
3133 if (result == NULL)
3134 goto empty;
3135 po->result = result;
3136 for (i=0; i<r ; i++) {
3137 index = indices[i];
3138 elem = PyTuple_GET_ITEM(pool, index);
3139 Py_INCREF(elem);
3140 PyTuple_SET_ITEM(result, i, elem);
3141 }
3142 } else {
3143 if (n == 0)
3144 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003146 /* Copy the previous result tuple or re-use it if available */
3147 if (Py_REFCNT(result) > 1) {
3148 PyObject *old_result = result;
3149 result = PyTuple_New(r);
3150 if (result == NULL)
3151 goto empty;
3152 po->result = result;
3153 for (i=0; i<r ; i++) {
3154 elem = PyTuple_GET_ITEM(old_result, i);
3155 Py_INCREF(elem);
3156 PyTuple_SET_ITEM(result, i, elem);
3157 }
3158 Py_DECREF(old_result);
3159 }
3160 /* Now, we've got the only copy so we can update it in-place */
3161 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003162
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003163 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3164 for (i=r-1 ; i>=0 ; i--) {
3165 cycles[i] -= 1;
3166 if (cycles[i] == 0) {
3167 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3168 index = indices[i];
3169 for (j=i ; j<n-1 ; j++)
3170 indices[j] = indices[j+1];
3171 indices[n-1] = index;
3172 cycles[i] = n - i;
3173 } else {
3174 j = cycles[i];
3175 index = indices[i];
3176 indices[i] = indices[n-j];
3177 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003178
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003179 for (k=i; k<r ; k++) {
3180 /* start with i, the leftmost element that changed */
3181 /* yield tuple(pool[k] for k in indices[:r]) */
3182 index = indices[k];
3183 elem = PyTuple_GET_ITEM(pool, index);
3184 Py_INCREF(elem);
3185 oldelem = PyTuple_GET_ITEM(result, k);
3186 PyTuple_SET_ITEM(result, k, elem);
3187 Py_DECREF(oldelem);
3188 }
3189 break;
3190 }
3191 }
3192 /* If i is negative, then the cycles have all
3193 rolled-over and we're done. */
3194 if (i < 0)
3195 goto empty;
3196 }
3197 Py_INCREF(result);
3198 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003199
3200empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003201 po->stopped = 1;
3202 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003203}
3204
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003205static PyObject *
3206permutations_reduce(permutationsobject *po)
3207{
3208 if (po->result == NULL) {
3209 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3210 } else if (po->stopped) {
3211 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3212 } else {
3213 PyObject *indices=NULL, *cycles=NULL;
3214 Py_ssize_t n, i;
3215
3216 /* we must pickle the indices and cycles and use them for setstate */
3217 n = PyTuple_GET_SIZE(po->pool);
3218 indices = PyTuple_New(n);
3219 if (indices == NULL)
3220 goto err;
3221 for (i=0; i<n; i++){
3222 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3223 if (!index)
3224 goto err;
3225 PyTuple_SET_ITEM(indices, i, index);
3226 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003227
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003228 cycles = PyTuple_New(po->r);
3229 if (cycles == NULL)
3230 goto err;
3231 for (i=0; i<po->r; i++)
3232 {
3233 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3234 if (!index)
3235 goto err;
3236 PyTuple_SET_ITEM(cycles, i, index);
3237 }
3238 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3239 po->pool, po->r,
3240 indices, cycles);
3241 err:
3242 Py_XDECREF(indices);
3243 Py_XDECREF(cycles);
3244 return NULL;
3245 }
3246}
3247
3248static PyObject *
3249permutations_setstate(permutationsobject *po, PyObject *state)
3250{
3251 PyObject *indices, *cycles, *result;
3252 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003253
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003254 if (!PyArg_ParseTuple(state, "O!O!",
3255 &PyTuple_Type, &indices,
3256 &PyTuple_Type, &cycles))
3257 return NULL;
3258
3259 n = PyTuple_GET_SIZE(po->pool);
3260 if (PyTuple_GET_SIZE(indices) != n ||
3261 PyTuple_GET_SIZE(cycles) != po->r)
3262 {
3263 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3264 return NULL;
3265 }
3266
3267 for (i=0; i<n; i++)
3268 {
3269 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3270 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3271 if (index < 0 && PyErr_Occurred())
3272 return NULL; /* not an integer */
3273 /* clamp the index */
3274 if (index < 0)
3275 index = 0;
3276 else if (index > n-1)
3277 index = n-1;
3278 po->indices[i] = index;
3279 }
3280
3281 for (i=0; i<po->r; i++)
3282 {
3283 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3284 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3285 if (index < 0 && PyErr_Occurred())
3286 return NULL; /* not an integer */
3287 if (index < 1)
3288 index = 1;
3289 else if (index > n-i)
3290 index = n-i;
3291 po->cycles[i] = index;
3292 }
3293 result = PyTuple_New(po->r);
3294 if (result == NULL)
3295 return NULL;
3296 for (i=0; i<po->r; i++) {
3297 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3298 Py_INCREF(element);
3299 PyTuple_SET_ITEM(result, i, element);
3300 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003301 Py_SETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003302 Py_RETURN_NONE;
3303}
3304
3305static PyMethodDef permuations_methods[] = {
3306 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3307 reduce_doc},
3308 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3309 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003310 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3311 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003312 {NULL, NULL} /* sentinel */
3313};
3314
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003315PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003316"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003317\n\
3318Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003319permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003320
3321static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003322 PyVarObject_HEAD_INIT(NULL, 0)
3323 "itertools.permutations", /* tp_name */
3324 sizeof(permutationsobject), /* tp_basicsize */
3325 0, /* tp_itemsize */
3326 /* methods */
3327 (destructor)permutations_dealloc, /* tp_dealloc */
3328 0, /* tp_print */
3329 0, /* tp_getattr */
3330 0, /* tp_setattr */
3331 0, /* tp_reserved */
3332 0, /* tp_repr */
3333 0, /* tp_as_number */
3334 0, /* tp_as_sequence */
3335 0, /* tp_as_mapping */
3336 0, /* tp_hash */
3337 0, /* tp_call */
3338 0, /* tp_str */
3339 PyObject_GenericGetAttr, /* tp_getattro */
3340 0, /* tp_setattro */
3341 0, /* tp_as_buffer */
3342 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3343 Py_TPFLAGS_BASETYPE, /* tp_flags */
3344 permutations_doc, /* tp_doc */
3345 (traverseproc)permutations_traverse, /* tp_traverse */
3346 0, /* tp_clear */
3347 0, /* tp_richcompare */
3348 0, /* tp_weaklistoffset */
3349 PyObject_SelfIter, /* tp_iter */
3350 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003351 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003352 0, /* tp_members */
3353 0, /* tp_getset */
3354 0, /* tp_base */
3355 0, /* tp_dict */
3356 0, /* tp_descr_get */
3357 0, /* tp_descr_set */
3358 0, /* tp_dictoffset */
3359 0, /* tp_init */
3360 0, /* tp_alloc */
3361 permutations_new, /* tp_new */
3362 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003363};
3364
Raymond Hettinger482ba772010-12-01 22:48:00 +00003365/* accumulate object ************************************************************/
3366
3367typedef struct {
3368 PyObject_HEAD
3369 PyObject *total;
3370 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003371 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003372} accumulateobject;
3373
3374static PyTypeObject accumulate_type;
3375
3376static PyObject *
3377accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3378{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003379 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003380 PyObject *iterable;
3381 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003382 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003383 accumulateobject *lz;
3384
Raymond Hettinger5d446132011-03-27 18:52:10 -07003385 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3386 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003387 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003388
3389 /* Get iterator. */
3390 it = PyObject_GetIter(iterable);
3391 if (it == NULL)
3392 return NULL;
3393
Raymond Hettinger482ba772010-12-01 22:48:00 +00003394 /* create accumulateobject structure */
3395 lz = (accumulateobject *)type->tp_alloc(type, 0);
3396 if (lz == NULL) {
3397 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003398 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003399 }
3400
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003401 if (binop != Py_None) {
3402 Py_XINCREF(binop);
3403 lz->binop = binop;
3404 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003405 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003406 lz->it = it;
3407 return (PyObject *)lz;
3408}
3409
3410static void
3411accumulate_dealloc(accumulateobject *lz)
3412{
3413 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003414 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003415 Py_XDECREF(lz->total);
3416 Py_XDECREF(lz->it);
3417 Py_TYPE(lz)->tp_free(lz);
3418}
3419
3420static int
3421accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3422{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003423 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003424 Py_VISIT(lz->it);
3425 Py_VISIT(lz->total);
3426 return 0;
3427}
3428
3429static PyObject *
3430accumulate_next(accumulateobject *lz)
3431{
3432 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003433
Raymond Hettinger482ba772010-12-01 22:48:00 +00003434 val = PyIter_Next(lz->it);
3435 if (val == NULL)
3436 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003437
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003438 if (lz->total == NULL) {
3439 Py_INCREF(val);
3440 lz->total = val;
3441 return lz->total;
3442 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003443
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003444 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003445 newtotal = PyNumber_Add(lz->total, val);
3446 else
3447 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003448 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003449 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003450 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003451
3452 oldtotal = lz->total;
3453 lz->total = newtotal;
3454 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003455
Raymond Hettinger482ba772010-12-01 22:48:00 +00003456 Py_INCREF(newtotal);
3457 return newtotal;
3458}
3459
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003460static PyObject *
3461accumulate_reduce(accumulateobject *lz)
3462{
3463 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3464 lz->it, lz->binop?lz->binop:Py_None,
3465 lz->total?lz->total:Py_None);
3466 }
3467
3468static PyObject *
3469accumulate_setstate(accumulateobject *lz, PyObject *state)
3470{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003471 Py_INCREF(state);
3472 Py_SETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003473 Py_RETURN_NONE;
3474}
3475
3476static PyMethodDef accumulate_methods[] = {
3477 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3478 reduce_doc},
3479 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3480 setstate_doc},
3481 {NULL, NULL} /* sentinel */
3482};
3483
Raymond Hettinger482ba772010-12-01 22:48:00 +00003484PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003485"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003486\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003487Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003488
3489static PyTypeObject accumulate_type = {
3490 PyVarObject_HEAD_INIT(NULL, 0)
3491 "itertools.accumulate", /* tp_name */
3492 sizeof(accumulateobject), /* tp_basicsize */
3493 0, /* tp_itemsize */
3494 /* methods */
3495 (destructor)accumulate_dealloc, /* tp_dealloc */
3496 0, /* tp_print */
3497 0, /* tp_getattr */
3498 0, /* tp_setattr */
3499 0, /* tp_reserved */
3500 0, /* tp_repr */
3501 0, /* tp_as_number */
3502 0, /* tp_as_sequence */
3503 0, /* tp_as_mapping */
3504 0, /* tp_hash */
3505 0, /* tp_call */
3506 0, /* tp_str */
3507 PyObject_GenericGetAttr, /* tp_getattro */
3508 0, /* tp_setattro */
3509 0, /* tp_as_buffer */
3510 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3511 Py_TPFLAGS_BASETYPE, /* tp_flags */
3512 accumulate_doc, /* tp_doc */
3513 (traverseproc)accumulate_traverse, /* tp_traverse */
3514 0, /* tp_clear */
3515 0, /* tp_richcompare */
3516 0, /* tp_weaklistoffset */
3517 PyObject_SelfIter, /* tp_iter */
3518 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003519 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003520 0, /* tp_members */
3521 0, /* tp_getset */
3522 0, /* tp_base */
3523 0, /* tp_dict */
3524 0, /* tp_descr_get */
3525 0, /* tp_descr_set */
3526 0, /* tp_dictoffset */
3527 0, /* tp_init */
3528 0, /* tp_alloc */
3529 accumulate_new, /* tp_new */
3530 PyObject_GC_Del, /* tp_free */
3531};
3532
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003533
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003534/* compress object ************************************************************/
3535
3536/* Equivalent to:
3537
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003538 def compress(data, selectors):
3539 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3540 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003541*/
3542
3543typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003544 PyObject_HEAD
3545 PyObject *data;
3546 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003547} compressobject;
3548
3549static PyTypeObject compress_type;
3550
3551static PyObject *
3552compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3553{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003554 PyObject *seq1, *seq2;
3555 PyObject *data=NULL, *selectors=NULL;
3556 compressobject *lz;
3557 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003558
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003559 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3560 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003561
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003562 data = PyObject_GetIter(seq1);
3563 if (data == NULL)
3564 goto fail;
3565 selectors = PyObject_GetIter(seq2);
3566 if (selectors == NULL)
3567 goto fail;
3568
3569 /* create compressobject structure */
3570 lz = (compressobject *)type->tp_alloc(type, 0);
3571 if (lz == NULL)
3572 goto fail;
3573 lz->data = data;
3574 lz->selectors = selectors;
3575 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003576
3577fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003578 Py_XDECREF(data);
3579 Py_XDECREF(selectors);
3580 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003581}
3582
3583static void
3584compress_dealloc(compressobject *lz)
3585{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003586 PyObject_GC_UnTrack(lz);
3587 Py_XDECREF(lz->data);
3588 Py_XDECREF(lz->selectors);
3589 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003590}
3591
3592static int
3593compress_traverse(compressobject *lz, visitproc visit, void *arg)
3594{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003595 Py_VISIT(lz->data);
3596 Py_VISIT(lz->selectors);
3597 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003598}
3599
3600static PyObject *
3601compress_next(compressobject *lz)
3602{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003603 PyObject *data = lz->data, *selectors = lz->selectors;
3604 PyObject *datum, *selector;
3605 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3606 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3607 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003608
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003609 while (1) {
3610 /* Steps: get datum, get selector, evaluate selector.
3611 Order is important (to match the pure python version
3612 in terms of which input gets a chance to raise an
3613 exception first).
3614 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003615
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003616 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003617 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003618 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003619
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003620 selector = selectornext(selectors);
3621 if (selector == NULL) {
3622 Py_DECREF(datum);
3623 return NULL;
3624 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003625
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003626 ok = PyObject_IsTrue(selector);
3627 Py_DECREF(selector);
3628 if (ok == 1)
3629 return datum;
3630 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003631 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003632 return NULL;
3633 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003634}
3635
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003636static PyObject *
3637compress_reduce(compressobject *lz)
3638{
3639 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3640 lz->data, lz->selectors);
3641 }
3642
3643static PyMethodDef compress_methods[] = {
3644 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3645 reduce_doc},
3646 {NULL, NULL} /* sentinel */
3647};
3648
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003649PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003650"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003651\n\
3652Return data elements corresponding to true selector elements.\n\
3653Forms a shorter iterator from selected data elements using the\n\
3654selectors to choose the data elements.");
3655
3656static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003657 PyVarObject_HEAD_INIT(NULL, 0)
3658 "itertools.compress", /* tp_name */
3659 sizeof(compressobject), /* tp_basicsize */
3660 0, /* tp_itemsize */
3661 /* methods */
3662 (destructor)compress_dealloc, /* tp_dealloc */
3663 0, /* tp_print */
3664 0, /* tp_getattr */
3665 0, /* tp_setattr */
3666 0, /* tp_reserved */
3667 0, /* tp_repr */
3668 0, /* tp_as_number */
3669 0, /* tp_as_sequence */
3670 0, /* tp_as_mapping */
3671 0, /* tp_hash */
3672 0, /* tp_call */
3673 0, /* tp_str */
3674 PyObject_GenericGetAttr, /* tp_getattro */
3675 0, /* tp_setattro */
3676 0, /* tp_as_buffer */
3677 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3678 Py_TPFLAGS_BASETYPE, /* tp_flags */
3679 compress_doc, /* tp_doc */
3680 (traverseproc)compress_traverse, /* tp_traverse */
3681 0, /* tp_clear */
3682 0, /* tp_richcompare */
3683 0, /* tp_weaklistoffset */
3684 PyObject_SelfIter, /* tp_iter */
3685 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003686 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003687 0, /* tp_members */
3688 0, /* tp_getset */
3689 0, /* tp_base */
3690 0, /* tp_dict */
3691 0, /* tp_descr_get */
3692 0, /* tp_descr_set */
3693 0, /* tp_dictoffset */
3694 0, /* tp_init */
3695 0, /* tp_alloc */
3696 compress_new, /* tp_new */
3697 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003698};
3699
3700
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003701/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003702
3703typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 PyObject_HEAD
3705 PyObject *func;
3706 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003707} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003708
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003709static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003710
3711static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003712filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003713{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003714 PyObject *func, *seq;
3715 PyObject *it;
3716 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003718 if (type == &filterfalse_type &&
3719 !_PyArg_NoKeywords("filterfalse()", kwds))
3720 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003721
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003722 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3723 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003724
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003725 /* Get iterator. */
3726 it = PyObject_GetIter(seq);
3727 if (it == NULL)
3728 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003729
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003730 /* create filterfalseobject structure */
3731 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3732 if (lz == NULL) {
3733 Py_DECREF(it);
3734 return NULL;
3735 }
3736 Py_INCREF(func);
3737 lz->func = func;
3738 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003739
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003740 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003741}
3742
3743static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003744filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003745{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003746 PyObject_GC_UnTrack(lz);
3747 Py_XDECREF(lz->func);
3748 Py_XDECREF(lz->it);
3749 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003750}
3751
3752static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003753filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003754{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003755 Py_VISIT(lz->it);
3756 Py_VISIT(lz->func);
3757 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003758}
3759
3760static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003761filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003762{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003763 PyObject *item;
3764 PyObject *it = lz->it;
3765 long ok;
3766 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003767
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003768 iternext = *Py_TYPE(it)->tp_iternext;
3769 for (;;) {
3770 item = iternext(it);
3771 if (item == NULL)
3772 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003773
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003774 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3775 ok = PyObject_IsTrue(item);
3776 } else {
3777 PyObject *good;
3778 good = PyObject_CallFunctionObjArgs(lz->func,
3779 item, NULL);
3780 if (good == NULL) {
3781 Py_DECREF(item);
3782 return NULL;
3783 }
3784 ok = PyObject_IsTrue(good);
3785 Py_DECREF(good);
3786 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003787 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003788 return item;
3789 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003790 if (ok < 0)
3791 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003792 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003793}
3794
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003795static PyObject *
3796filterfalse_reduce(filterfalseobject *lz)
3797{
3798 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3799 lz->func, lz->it);
3800 }
3801
3802static PyMethodDef filterfalse_methods[] = {
3803 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3804 reduce_doc},
3805 {NULL, NULL} /* sentinel */
3806};
3807
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003808PyDoc_STRVAR(filterfalse_doc,
3809"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003810\n\
3811Return those items of sequence for which function(item) is false.\n\
3812If function is None, return the items that are false.");
3813
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003814static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003815 PyVarObject_HEAD_INIT(NULL, 0)
3816 "itertools.filterfalse", /* tp_name */
3817 sizeof(filterfalseobject), /* tp_basicsize */
3818 0, /* tp_itemsize */
3819 /* methods */
3820 (destructor)filterfalse_dealloc, /* tp_dealloc */
3821 0, /* tp_print */
3822 0, /* tp_getattr */
3823 0, /* tp_setattr */
3824 0, /* tp_reserved */
3825 0, /* tp_repr */
3826 0, /* tp_as_number */
3827 0, /* tp_as_sequence */
3828 0, /* tp_as_mapping */
3829 0, /* tp_hash */
3830 0, /* tp_call */
3831 0, /* tp_str */
3832 PyObject_GenericGetAttr, /* tp_getattro */
3833 0, /* tp_setattro */
3834 0, /* tp_as_buffer */
3835 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3836 Py_TPFLAGS_BASETYPE, /* tp_flags */
3837 filterfalse_doc, /* tp_doc */
3838 (traverseproc)filterfalse_traverse, /* tp_traverse */
3839 0, /* tp_clear */
3840 0, /* tp_richcompare */
3841 0, /* tp_weaklistoffset */
3842 PyObject_SelfIter, /* tp_iter */
3843 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003844 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003845 0, /* tp_members */
3846 0, /* tp_getset */
3847 0, /* tp_base */
3848 0, /* tp_dict */
3849 0, /* tp_descr_get */
3850 0, /* tp_descr_set */
3851 0, /* tp_dictoffset */
3852 0, /* tp_init */
3853 0, /* tp_alloc */
3854 filterfalse_new, /* tp_new */
3855 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003856};
3857
3858
3859/* count object ************************************************************/
3860
3861typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003862 PyObject_HEAD
3863 Py_ssize_t cnt;
3864 PyObject *long_cnt;
3865 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003866} countobject;
3867
Raymond Hettinger30729212009-02-12 06:28:27 +00003868/* Counting logic and invariants:
3869
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003870fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003871
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003872 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003873 Advances with: cnt += 1
3874 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003875
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003876slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003878 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3879 All counting is done with python objects (no overflows or underflows).
3880 Advances with: long_cnt += long_step
3881 Step may be zero -- effectively a slow version of repeat(cnt).
3882 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003883*/
3884
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003885static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003886
3887static PyObject *
3888count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3889{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003890 countobject *lz;
3891 int slow_mode = 0;
3892 Py_ssize_t cnt = 0;
3893 PyObject *long_cnt = NULL;
3894 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003895 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003896 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003897
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003898 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3899 kwlist, &long_cnt, &long_step))
3900 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003901
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003902 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3903 (long_step != NULL && !PyNumber_Check(long_step))) {
3904 PyErr_SetString(PyExc_TypeError, "a number is required");
3905 return NULL;
3906 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003907
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003908 if (long_cnt != NULL) {
3909 cnt = PyLong_AsSsize_t(long_cnt);
3910 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3911 PyErr_Clear();
3912 slow_mode = 1;
3913 }
3914 Py_INCREF(long_cnt);
3915 } else {
3916 cnt = 0;
3917 long_cnt = PyLong_FromLong(0);
3918 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003919
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003920 /* If not specified, step defaults to 1 */
3921 if (long_step == NULL) {
3922 long_step = PyLong_FromLong(1);
3923 if (long_step == NULL) {
3924 Py_DECREF(long_cnt);
3925 return NULL;
3926 }
3927 } else
3928 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003929
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003930 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003931
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003932 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003933 step = PyLong_AsLong(long_step);
3934 if (step != 1) {
3935 slow_mode = 1;
3936 if (step == -1 && PyErr_Occurred())
3937 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003938 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003939
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003940 if (slow_mode)
3941 cnt = PY_SSIZE_T_MAX;
3942 else
3943 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003944
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3946 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3947 assert(slow_mode ||
3948 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003949
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003950 /* create countobject structure */
3951 lz = (countobject *)type->tp_alloc(type, 0);
3952 if (lz == NULL) {
3953 Py_XDECREF(long_cnt);
3954 return NULL;
3955 }
3956 lz->cnt = cnt;
3957 lz->long_cnt = long_cnt;
3958 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003959
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003960 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003961}
3962
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003963static void
3964count_dealloc(countobject *lz)
3965{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003966 PyObject_GC_UnTrack(lz);
3967 Py_XDECREF(lz->long_cnt);
3968 Py_XDECREF(lz->long_step);
3969 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003970}
3971
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003972static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003973count_traverse(countobject *lz, visitproc visit, void *arg)
3974{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 Py_VISIT(lz->long_cnt);
3976 Py_VISIT(lz->long_step);
3977 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003978}
3979
3980static PyObject *
3981count_nextlong(countobject *lz)
3982{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003983 PyObject *long_cnt;
3984 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003985
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003986 long_cnt = lz->long_cnt;
3987 if (long_cnt == NULL) {
3988 /* Switch to slow_mode */
3989 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3990 if (long_cnt == NULL)
3991 return NULL;
3992 }
3993 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003994
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003995 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3996 if (stepped_up == NULL)
3997 return NULL;
3998 lz->long_cnt = stepped_up;
3999 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004000}
4001
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004002static PyObject *
4003count_next(countobject *lz)
4004{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004005 if (lz->cnt == PY_SSIZE_T_MAX)
4006 return count_nextlong(lz);
4007 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004008}
4009
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004010static PyObject *
4011count_repr(countobject *lz)
4012{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004013 if (lz->cnt != PY_SSIZE_T_MAX)
4014 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004015
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004016 if (PyLong_Check(lz->long_step)) {
4017 long step = PyLong_AsLong(lz->long_step);
4018 if (step == -1 && PyErr_Occurred()) {
4019 PyErr_Clear();
4020 }
4021 if (step == 1) {
4022 /* Don't display step when it is an integer equal to 1 */
4023 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4024 }
4025 }
4026 return PyUnicode_FromFormat("count(%R, %R)",
4027 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004028}
4029
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004030static PyObject *
4031count_reduce(countobject *lz)
4032{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004033 if (lz->cnt == PY_SSIZE_T_MAX)
4034 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4035 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004036}
4037
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004038static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004039 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004040 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004041 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004042};
4043
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004044PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004045 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004046\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004047Return a count object whose .__next__() method returns consecutive values.\n\
4048Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004049 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004050 x = firstval\n\
4051 while 1:\n\
4052 yield x\n\
4053 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004054
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004055static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004056 PyVarObject_HEAD_INIT(NULL, 0)
4057 "itertools.count", /* tp_name */
4058 sizeof(countobject), /* tp_basicsize */
4059 0, /* tp_itemsize */
4060 /* methods */
4061 (destructor)count_dealloc, /* tp_dealloc */
4062 0, /* tp_print */
4063 0, /* tp_getattr */
4064 0, /* tp_setattr */
4065 0, /* tp_reserved */
4066 (reprfunc)count_repr, /* tp_repr */
4067 0, /* tp_as_number */
4068 0, /* tp_as_sequence */
4069 0, /* tp_as_mapping */
4070 0, /* tp_hash */
4071 0, /* tp_call */
4072 0, /* tp_str */
4073 PyObject_GenericGetAttr, /* tp_getattro */
4074 0, /* tp_setattro */
4075 0, /* tp_as_buffer */
4076 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4077 Py_TPFLAGS_BASETYPE, /* tp_flags */
4078 count_doc, /* tp_doc */
4079 (traverseproc)count_traverse, /* tp_traverse */
4080 0, /* tp_clear */
4081 0, /* tp_richcompare */
4082 0, /* tp_weaklistoffset */
4083 PyObject_SelfIter, /* tp_iter */
4084 (iternextfunc)count_next, /* tp_iternext */
4085 count_methods, /* tp_methods */
4086 0, /* tp_members */
4087 0, /* tp_getset */
4088 0, /* tp_base */
4089 0, /* tp_dict */
4090 0, /* tp_descr_get */
4091 0, /* tp_descr_set */
4092 0, /* tp_dictoffset */
4093 0, /* tp_init */
4094 0, /* tp_alloc */
4095 count_new, /* tp_new */
4096 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004097};
4098
4099
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004100/* repeat object ************************************************************/
4101
4102typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004103 PyObject_HEAD
4104 PyObject *element;
4105 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004106} repeatobject;
4107
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004108static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004109
4110static PyObject *
4111repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004113 repeatobject *ro;
4114 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004115 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004116 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004118 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4119 &element, &cnt))
4120 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004121
Raymond Hettinger97d35552014-06-24 21:36:58 -07004122 if (kwds != NULL)
4123 n_kwds = PyDict_Size(kwds);
4124 /* Does user supply times argument? */
4125 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004126 cnt = 0;
4127
4128 ro = (repeatobject *)type->tp_alloc(type, 0);
4129 if (ro == NULL)
4130 return NULL;
4131 Py_INCREF(element);
4132 ro->element = element;
4133 ro->cnt = cnt;
4134 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004135}
4136
4137static void
4138repeat_dealloc(repeatobject *ro)
4139{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004140 PyObject_GC_UnTrack(ro);
4141 Py_XDECREF(ro->element);
4142 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004143}
4144
4145static int
4146repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4147{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004148 Py_VISIT(ro->element);
4149 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004150}
4151
4152static PyObject *
4153repeat_next(repeatobject *ro)
4154{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004155 if (ro->cnt == 0)
4156 return NULL;
4157 if (ro->cnt > 0)
4158 ro->cnt--;
4159 Py_INCREF(ro->element);
4160 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004161}
4162
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004163static PyObject *
4164repeat_repr(repeatobject *ro)
4165{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004166 if (ro->cnt == -1)
4167 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4168 else
4169 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4170}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004171
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004172static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004173repeat_len(repeatobject *ro)
4174{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004175 if (ro->cnt == -1) {
4176 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4177 return NULL;
4178 }
4179 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004180}
4181
Armin Rigof5b3e362006-02-11 21:32:43 +00004182PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004183
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004184static PyObject *
4185repeat_reduce(repeatobject *ro)
4186{
4187 /* unpickle this so that a new repeat iterator is constructed with an
4188 * object, then call __setstate__ on it to set cnt
4189 */
4190 if (ro->cnt >= 0)
4191 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4192 else
4193 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4194}
4195
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004196static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004197 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004198 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004199 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004200};
4201
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004202PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004203"repeat(object [,times]) -> create an iterator which returns the object\n\
4204for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004205endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004206
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004207static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004208 PyVarObject_HEAD_INIT(NULL, 0)
4209 "itertools.repeat", /* tp_name */
4210 sizeof(repeatobject), /* tp_basicsize */
4211 0, /* tp_itemsize */
4212 /* methods */
4213 (destructor)repeat_dealloc, /* tp_dealloc */
4214 0, /* tp_print */
4215 0, /* tp_getattr */
4216 0, /* tp_setattr */
4217 0, /* tp_reserved */
4218 (reprfunc)repeat_repr, /* tp_repr */
4219 0, /* tp_as_number */
4220 0, /* tp_as_sequence */
4221 0, /* tp_as_mapping */
4222 0, /* tp_hash */
4223 0, /* tp_call */
4224 0, /* tp_str */
4225 PyObject_GenericGetAttr, /* tp_getattro */
4226 0, /* tp_setattro */
4227 0, /* tp_as_buffer */
4228 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4229 Py_TPFLAGS_BASETYPE, /* tp_flags */
4230 repeat_doc, /* tp_doc */
4231 (traverseproc)repeat_traverse, /* tp_traverse */
4232 0, /* tp_clear */
4233 0, /* tp_richcompare */
4234 0, /* tp_weaklistoffset */
4235 PyObject_SelfIter, /* tp_iter */
4236 (iternextfunc)repeat_next, /* tp_iternext */
4237 repeat_methods, /* tp_methods */
4238 0, /* tp_members */
4239 0, /* tp_getset */
4240 0, /* tp_base */
4241 0, /* tp_dict */
4242 0, /* tp_descr_get */
4243 0, /* tp_descr_set */
4244 0, /* tp_dictoffset */
4245 0, /* tp_init */
4246 0, /* tp_alloc */
4247 repeat_new, /* tp_new */
4248 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004249};
4250
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004251/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004252
4253#include "Python.h"
4254
4255typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004256 PyObject_HEAD
4257 Py_ssize_t tuplesize;
4258 Py_ssize_t numactive;
4259 PyObject *ittuple; /* tuple of iterators */
4260 PyObject *result;
4261 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004262} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004263
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004264static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004265
4266static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004267zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004268{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004269 ziplongestobject *lz;
4270 Py_ssize_t i;
4271 PyObject *ittuple; /* tuple of iterators */
4272 PyObject *result;
4273 PyObject *fillvalue = Py_None;
4274 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004275
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004276 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4277 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4278 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4279 PyErr_SetString(PyExc_TypeError,
4280 "zip_longest() got an unexpected keyword argument");
4281 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004282 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004283 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004284
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004285 /* args must be a tuple */
4286 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004287
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004288 /* obtain iterators */
4289 ittuple = PyTuple_New(tuplesize);
4290 if (ittuple == NULL)
4291 return NULL;
4292 for (i=0; i < tuplesize; ++i) {
4293 PyObject *item = PyTuple_GET_ITEM(args, i);
4294 PyObject *it = PyObject_GetIter(item);
4295 if (it == NULL) {
4296 if (PyErr_ExceptionMatches(PyExc_TypeError))
4297 PyErr_Format(PyExc_TypeError,
4298 "zip_longest argument #%zd must support iteration",
4299 i+1);
4300 Py_DECREF(ittuple);
4301 return NULL;
4302 }
4303 PyTuple_SET_ITEM(ittuple, i, it);
4304 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004305
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004306 /* create a result holder */
4307 result = PyTuple_New(tuplesize);
4308 if (result == NULL) {
4309 Py_DECREF(ittuple);
4310 return NULL;
4311 }
4312 for (i=0 ; i < tuplesize ; i++) {
4313 Py_INCREF(Py_None);
4314 PyTuple_SET_ITEM(result, i, Py_None);
4315 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004317 /* create ziplongestobject structure */
4318 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4319 if (lz == NULL) {
4320 Py_DECREF(ittuple);
4321 Py_DECREF(result);
4322 return NULL;
4323 }
4324 lz->ittuple = ittuple;
4325 lz->tuplesize = tuplesize;
4326 lz->numactive = tuplesize;
4327 lz->result = result;
4328 Py_INCREF(fillvalue);
4329 lz->fillvalue = fillvalue;
4330 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004331}
4332
4333static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004334zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004335{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004336 PyObject_GC_UnTrack(lz);
4337 Py_XDECREF(lz->ittuple);
4338 Py_XDECREF(lz->result);
4339 Py_XDECREF(lz->fillvalue);
4340 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004341}
4342
4343static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004344zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004345{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004346 Py_VISIT(lz->ittuple);
4347 Py_VISIT(lz->result);
4348 Py_VISIT(lz->fillvalue);
4349 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004350}
4351
4352static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004353zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004354{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004355 Py_ssize_t i;
4356 Py_ssize_t tuplesize = lz->tuplesize;
4357 PyObject *result = lz->result;
4358 PyObject *it;
4359 PyObject *item;
4360 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004361
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004362 if (tuplesize == 0)
4363 return NULL;
4364 if (lz->numactive == 0)
4365 return NULL;
4366 if (Py_REFCNT(result) == 1) {
4367 Py_INCREF(result);
4368 for (i=0 ; i < tuplesize ; i++) {
4369 it = PyTuple_GET_ITEM(lz->ittuple, i);
4370 if (it == NULL) {
4371 Py_INCREF(lz->fillvalue);
4372 item = lz->fillvalue;
4373 } else {
4374 item = PyIter_Next(it);
4375 if (item == NULL) {
4376 lz->numactive -= 1;
4377 if (lz->numactive == 0 || PyErr_Occurred()) {
4378 lz->numactive = 0;
4379 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004380 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004381 } else {
4382 Py_INCREF(lz->fillvalue);
4383 item = lz->fillvalue;
4384 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4385 Py_DECREF(it);
4386 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004387 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004388 }
4389 olditem = PyTuple_GET_ITEM(result, i);
4390 PyTuple_SET_ITEM(result, i, item);
4391 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004392 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004393 } else {
4394 result = PyTuple_New(tuplesize);
4395 if (result == NULL)
4396 return NULL;
4397 for (i=0 ; i < tuplesize ; i++) {
4398 it = PyTuple_GET_ITEM(lz->ittuple, i);
4399 if (it == NULL) {
4400 Py_INCREF(lz->fillvalue);
4401 item = lz->fillvalue;
4402 } else {
4403 item = PyIter_Next(it);
4404 if (item == NULL) {
4405 lz->numactive -= 1;
4406 if (lz->numactive == 0 || PyErr_Occurred()) {
4407 lz->numactive = 0;
4408 Py_DECREF(result);
4409 return NULL;
4410 } else {
4411 Py_INCREF(lz->fillvalue);
4412 item = lz->fillvalue;
4413 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4414 Py_DECREF(it);
4415 }
4416 }
4417 }
4418 PyTuple_SET_ITEM(result, i, item);
4419 }
4420 }
4421 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004422}
4423
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004424static PyObject *
4425zip_longest_reduce(ziplongestobject *lz)
4426{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004427
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004428 /* Create a new tuple with empty sequences where appropriate to pickle.
4429 * Then use setstate to set the fillvalue
4430 */
4431 int i;
4432 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4433 if (args == NULL)
4434 return NULL;
4435 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4436 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4437 if (elem == NULL) {
4438 elem = PyTuple_New(0);
4439 if (elem == NULL) {
4440 Py_DECREF(args);
4441 return NULL;
4442 }
4443 } else
4444 Py_INCREF(elem);
4445 PyTuple_SET_ITEM(args, i, elem);
4446 }
4447 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4448}
4449
4450static PyObject *
4451zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4452{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004453 Py_INCREF(state);
4454 Py_SETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004455 Py_RETURN_NONE;
4456}
4457
4458static PyMethodDef zip_longest_methods[] = {
4459 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4460 reduce_doc},
4461 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4462 setstate_doc},
4463 {NULL, NULL} /* sentinel */
4464};
4465
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004466PyDoc_STRVAR(zip_longest_doc,
4467"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004468\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004469Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004470the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004471method continues until the longest iterable in the argument sequence\n\
4472is exhausted and then it raises StopIteration. When the shorter iterables\n\
4473are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4474defaults to None or can be specified by a keyword argument.\n\
4475");
4476
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004477static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004478 PyVarObject_HEAD_INIT(NULL, 0)
4479 "itertools.zip_longest", /* tp_name */
4480 sizeof(ziplongestobject), /* tp_basicsize */
4481 0, /* tp_itemsize */
4482 /* methods */
4483 (destructor)zip_longest_dealloc, /* tp_dealloc */
4484 0, /* tp_print */
4485 0, /* tp_getattr */
4486 0, /* tp_setattr */
4487 0, /* tp_reserved */
4488 0, /* tp_repr */
4489 0, /* tp_as_number */
4490 0, /* tp_as_sequence */
4491 0, /* tp_as_mapping */
4492 0, /* tp_hash */
4493 0, /* tp_call */
4494 0, /* tp_str */
4495 PyObject_GenericGetAttr, /* tp_getattro */
4496 0, /* tp_setattro */
4497 0, /* tp_as_buffer */
4498 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4499 Py_TPFLAGS_BASETYPE, /* tp_flags */
4500 zip_longest_doc, /* tp_doc */
4501 (traverseproc)zip_longest_traverse, /* tp_traverse */
4502 0, /* tp_clear */
4503 0, /* tp_richcompare */
4504 0, /* tp_weaklistoffset */
4505 PyObject_SelfIter, /* tp_iter */
4506 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004507 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004508 0, /* tp_members */
4509 0, /* tp_getset */
4510 0, /* tp_base */
4511 0, /* tp_dict */
4512 0, /* tp_descr_get */
4513 0, /* tp_descr_set */
4514 0, /* tp_dictoffset */
4515 0, /* tp_init */
4516 0, /* tp_alloc */
4517 zip_longest_new, /* tp_new */
4518 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004519};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004520
4521/* module level code ********************************************************/
4522
4523PyDoc_STRVAR(module_doc,
4524"Functional tools for creating and using iterators.\n\
4525\n\
4526Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004527count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004528cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004529repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004530\n\
4531Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004532accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004533chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004534chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004535compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4536dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4537groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004538filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004539islice(seq, [start,] stop [, step]) --> elements from\n\
4540 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004541starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004542tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004543takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004544zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004545\n\
4546Combinatoric generators:\n\
4547product(p, q, ... [repeat=1]) --> cartesian product\n\
4548permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004549combinations(p, r)\n\
4550combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004551");
4552
4553
Raymond Hettingerad983e72003-11-12 14:32:26 +00004554static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004555 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4556 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004557};
4558
Martin v. Löwis1a214512008-06-11 05:26:20 +00004559
4560static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004561 PyModuleDef_HEAD_INIT,
4562 "itertools",
4563 module_doc,
4564 -1,
4565 module_methods,
4566 NULL,
4567 NULL,
4568 NULL,
4569 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004570};
4571
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004572PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004573PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004574{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004575 int i;
4576 PyObject *m;
4577 char *name;
4578 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004579 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004580 &combinations_type,
4581 &cwr_type,
4582 &cycle_type,
4583 &dropwhile_type,
4584 &takewhile_type,
4585 &islice_type,
4586 &starmap_type,
4587 &chain_type,
4588 &compress_type,
4589 &filterfalse_type,
4590 &count_type,
4591 &ziplongest_type,
4592 &permutations_type,
4593 &product_type,
4594 &repeat_type,
4595 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004596 &_grouper_type,
4597 &tee_type,
4598 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004599 NULL
4600 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004601
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004602 Py_TYPE(&teedataobject_type) = &PyType_Type;
4603 m = PyModule_Create(&itertoolsmodule);
4604 if (m == NULL)
4605 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004606
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004607 for (i=0 ; typelist[i] != NULL ; i++) {
4608 if (PyType_Ready(typelist[i]) < 0)
4609 return NULL;
4610 name = strchr(typelist[i]->tp_name, '.');
4611 assert (name != NULL);
4612 Py_INCREF(typelist[i]);
4613 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4614 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004615
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004616 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004617}