blob: ae5b166561d081e6ee58a563e82207ab4852628e [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 Hettinger8df58f72013-09-09 01:29:40 -05007 Copyright (c) 2003-2013 Python Software Foundation.
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000137static PyObject *
138groupby_reduce(groupbyobject *lz)
139{
140 /* reduce as a 'new' call with an optional 'setstate' if groupby
141 * has started
142 */
143 PyObject *value;
144 if (lz->tgtkey && lz->currkey && lz->currvalue)
145 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
146 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
147 else
148 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
149 lz->it, lz->keyfunc);
150
151 return value;
152}
153
154PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
155
156static PyObject *
157groupby_setstate(groupbyobject *lz, PyObject *state)
158{
159 PyObject *currkey, *currvalue, *tgtkey;
160 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
161 return NULL;
162 Py_CLEAR(lz->currkey);
163 lz->currkey = currkey;
164 Py_INCREF(lz->currkey);
165 Py_CLEAR(lz->currvalue);
166 lz->currvalue = currvalue;
167 Py_INCREF(lz->currvalue);
168 Py_CLEAR(lz->tgtkey);
169 lz->tgtkey = tgtkey;
170 Py_INCREF(lz->tgtkey);
171 Py_RETURN_NONE;
172}
173
174PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
175
176static PyMethodDef groupby_methods[] = {
177 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
178 reduce_doc},
179 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
180 setstate_doc},
181 {NULL, NULL} /* sentinel */
182};
183
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000184PyDoc_STRVAR(groupby_doc,
185"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
186(key, sub-iterator) grouped by each value of key(value).\n");
187
188static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyVarObject_HEAD_INIT(NULL, 0)
190 "itertools.groupby", /* tp_name */
191 sizeof(groupbyobject), /* tp_basicsize */
192 0, /* tp_itemsize */
193 /* methods */
194 (destructor)groupby_dealloc, /* tp_dealloc */
195 0, /* tp_print */
196 0, /* tp_getattr */
197 0, /* tp_setattr */
198 0, /* tp_reserved */
199 0, /* tp_repr */
200 0, /* tp_as_number */
201 0, /* tp_as_sequence */
202 0, /* tp_as_mapping */
203 0, /* tp_hash */
204 0, /* tp_call */
205 0, /* tp_str */
206 PyObject_GenericGetAttr, /* tp_getattro */
207 0, /* tp_setattro */
208 0, /* tp_as_buffer */
209 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
210 Py_TPFLAGS_BASETYPE, /* tp_flags */
211 groupby_doc, /* tp_doc */
212 (traverseproc)groupby_traverse, /* tp_traverse */
213 0, /* tp_clear */
214 0, /* tp_richcompare */
215 0, /* tp_weaklistoffset */
216 PyObject_SelfIter, /* tp_iter */
217 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000218 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000219 0, /* tp_members */
220 0, /* tp_getset */
221 0, /* tp_base */
222 0, /* tp_dict */
223 0, /* tp_descr_get */
224 0, /* tp_descr_set */
225 0, /* tp_dictoffset */
226 0, /* tp_init */
227 0, /* tp_alloc */
228 groupby_new, /* tp_new */
229 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000230};
231
232
233/* _grouper object (internal) ************************************************/
234
235typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000236 PyObject_HEAD
237 PyObject *parent;
238 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000239} _grouperobject;
240
241static PyTypeObject _grouper_type;
242
243static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000244_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
245{
246 PyObject *parent, *tgtkey;
247
248 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
249 return NULL;
250
251 return _grouper_create((groupbyobject*) parent, tgtkey);
252}
253
254static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000255_grouper_create(groupbyobject *parent, PyObject *tgtkey)
256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000257 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
260 if (igo == NULL)
261 return NULL;
262 igo->parent = (PyObject *)parent;
263 Py_INCREF(parent);
264 igo->tgtkey = tgtkey;
265 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000267 PyObject_GC_Track(igo);
268 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269}
270
271static void
272_grouper_dealloc(_grouperobject *igo)
273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000274 PyObject_GC_UnTrack(igo);
275 Py_DECREF(igo->parent);
276 Py_DECREF(igo->tgtkey);
277 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000278}
279
280static int
281_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000283 Py_VISIT(igo->parent);
284 Py_VISIT(igo->tgtkey);
285 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286}
287
288static PyObject *
289_grouper_next(_grouperobject *igo)
290{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000291 groupbyobject *gbo = (groupbyobject *)igo->parent;
292 PyObject *newvalue, *newkey, *r;
293 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (gbo->currvalue == NULL) {
296 newvalue = PyIter_Next(gbo->it);
297 if (newvalue == NULL)
298 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 if (gbo->keyfunc == Py_None) {
301 newkey = newvalue;
302 Py_INCREF(newvalue);
303 } else {
304 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
305 newvalue, NULL);
306 if (newkey == NULL) {
307 Py_DECREF(newvalue);
308 return NULL;
309 }
310 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 assert(gbo->currkey == NULL);
313 gbo->currkey = newkey;
314 gbo->currvalue = newvalue;
315 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000317 assert(gbo->currkey != NULL);
318 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
319 if (rcmp <= 0)
320 /* got any error or current group is end */
321 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000323 r = gbo->currvalue;
324 gbo->currvalue = NULL;
325 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000326
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000327 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000328}
329
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000330static PyObject *
331_grouper_reduce(_grouperobject *lz)
332{
333 return Py_BuildValue("O(OO)", Py_TYPE(lz),
334 lz->parent, lz->tgtkey);
335}
336
337static PyMethodDef _grouper_methods[] = {
338 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
339 reduce_doc},
340 {NULL, NULL} /* sentinel */
341};
342
343
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000344static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000345 PyVarObject_HEAD_INIT(NULL, 0)
346 "itertools._grouper", /* tp_name */
347 sizeof(_grouperobject), /* tp_basicsize */
348 0, /* tp_itemsize */
349 /* methods */
350 (destructor)_grouper_dealloc, /* tp_dealloc */
351 0, /* tp_print */
352 0, /* tp_getattr */
353 0, /* tp_setattr */
354 0, /* tp_reserved */
355 0, /* tp_repr */
356 0, /* tp_as_number */
357 0, /* tp_as_sequence */
358 0, /* tp_as_mapping */
359 0, /* tp_hash */
360 0, /* tp_call */
361 0, /* tp_str */
362 PyObject_GenericGetAttr, /* tp_getattro */
363 0, /* tp_setattro */
364 0, /* tp_as_buffer */
365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
366 0, /* tp_doc */
367 (traverseproc)_grouper_traverse,/* tp_traverse */
368 0, /* tp_clear */
369 0, /* tp_richcompare */
370 0, /* tp_weaklistoffset */
371 PyObject_SelfIter, /* tp_iter */
372 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000373 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000374 0, /* tp_members */
375 0, /* tp_getset */
376 0, /* tp_base */
377 0, /* tp_dict */
378 0, /* tp_descr_get */
379 0, /* tp_descr_set */
380 0, /* tp_dictoffset */
381 0, /* tp_init */
382 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000383 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000384 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000385};
386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000387
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000390
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391/* The teedataobject pre-allocates space for LINKCELLS number of objects.
392 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000393 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394 the other structure members including PyHEAD overhead. The larger the
395 value, the less memory overhead per object and the less time spent
396 allocating/deallocating new links. The smaller the number, the less
397 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400
401typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000402 PyObject_HEAD
403 PyObject *it;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200404 int numread; /* 0 <= numread <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000405 PyObject *nextlink;
406 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000407} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000408
409typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 PyObject_HEAD
411 teedataobject *dataobj;
Antoine Pitroub4a46cb2013-09-20 22:19:22 +0200412 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000413 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417
418static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000419teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000421 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000423 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
424 if (tdo == NULL)
425 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427 tdo->numread = 0;
428 tdo->nextlink = NULL;
429 Py_INCREF(it);
430 tdo->it = it;
431 PyObject_GC_Track(tdo);
432 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433}
434
435static PyObject *
436teedataobject_jumplink(teedataobject *tdo)
437{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000439 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000440 Py_XINCREF(tdo->nextlink);
441 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442}
443
444static PyObject *
445teedataobject_getitem(teedataobject *tdo, int i)
446{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000449 assert(i < LINKCELLS);
450 if (i < tdo->numread)
451 value = tdo->values[i];
452 else {
453 /* this is the lead iterator, so fetch more data */
454 assert(i == tdo->numread);
455 value = PyIter_Next(tdo->it);
456 if (value == NULL)
457 return NULL;
458 tdo->numread++;
459 tdo->values[i] = value;
460 }
461 Py_INCREF(value);
462 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463}
464
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000465static int
466teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
467{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 int i;
469 Py_VISIT(tdo->it);
470 for (i = 0; i < tdo->numread; i++)
471 Py_VISIT(tdo->values[i]);
472 Py_VISIT(tdo->nextlink);
473 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000474}
475
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200476static void
477teedataobject_safe_decref(PyObject *obj)
478{
479 while (obj && Py_TYPE(obj) == &teedataobject_type &&
480 Py_REFCNT(obj) == 1) {
481 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
482 ((teedataobject *)obj)->nextlink = NULL;
483 Py_DECREF(obj);
484 obj = nextlink;
485 }
486 Py_XDECREF(obj);
487}
488
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000489static int
490teedataobject_clear(teedataobject *tdo)
491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200493 PyObject *tmp;
494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000495 Py_CLEAR(tdo->it);
496 for (i=0 ; i<tdo->numread ; i++)
497 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200498 tmp = tdo->nextlink;
499 tdo->nextlink = NULL;
500 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000501 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000502}
503
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000505teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000507 PyObject_GC_UnTrack(tdo);
508 teedataobject_clear(tdo);
509 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000510}
511
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000512static PyObject *
513teedataobject_reduce(teedataobject *tdo)
514{
515 int i;
516 /* create a temporary list of already iterated values */
517 PyObject *values = PyList_New(tdo->numread);
518 if (!values)
519 return NULL;
520 for (i=0 ; i<tdo->numread ; i++) {
521 Py_INCREF(tdo->values[i]);
522 PyList_SET_ITEM(values, i, tdo->values[i]);
523 }
524 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
525 values,
526 tdo->nextlink ? tdo->nextlink : Py_None);
527}
528
529static PyTypeObject teedataobject_type;
530
531static PyObject *
532teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
533{
534 teedataobject *tdo;
535 PyObject *it, *values, *next;
536 Py_ssize_t i, len;
537
538 assert(type == &teedataobject_type);
539 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
540 return NULL;
541
542 tdo = (teedataobject *)teedataobject_newinternal(it);
543 if (!tdo)
544 return NULL;
545
546 len = PyList_GET_SIZE(values);
547 if (len > LINKCELLS)
548 goto err;
549 for (i=0; i<len; i++) {
550 tdo->values[i] = PyList_GET_ITEM(values, i);
551 Py_INCREF(tdo->values[i]);
552 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200553 /* len <= LINKCELLS < INT_MAX */
554 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000555
556 if (len == LINKCELLS) {
557 if (next != Py_None) {
558 if (Py_TYPE(next) != &teedataobject_type)
559 goto err;
560 assert(tdo->nextlink == NULL);
561 Py_INCREF(next);
562 tdo->nextlink = next;
563 }
564 } else {
565 if (next != Py_None)
566 goto err; /* shouldn't have a next if we are not full */
567 }
568 return (PyObject*)tdo;
569
570err:
571 Py_XDECREF(tdo);
572 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
573 return NULL;
574}
575
576static PyMethodDef teedataobject_methods[] = {
577 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
578 reduce_doc},
579 {NULL, NULL} /* sentinel */
580};
581
Raymond Hettingerad983e72003-11-12 14:32:26 +0000582PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583
Raymond Hettingerad983e72003-11-12 14:32:26 +0000584static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000585 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000586 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000587 sizeof(teedataobject), /* tp_basicsize */
588 0, /* tp_itemsize */
589 /* methods */
590 (destructor)teedataobject_dealloc, /* tp_dealloc */
591 0, /* tp_print */
592 0, /* tp_getattr */
593 0, /* tp_setattr */
594 0, /* tp_reserved */
595 0, /* tp_repr */
596 0, /* tp_as_number */
597 0, /* tp_as_sequence */
598 0, /* tp_as_mapping */
599 0, /* tp_hash */
600 0, /* tp_call */
601 0, /* tp_str */
602 PyObject_GenericGetAttr, /* tp_getattro */
603 0, /* tp_setattro */
604 0, /* tp_as_buffer */
605 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
606 teedataobject_doc, /* tp_doc */
607 (traverseproc)teedataobject_traverse, /* tp_traverse */
608 (inquiry)teedataobject_clear, /* tp_clear */
609 0, /* tp_richcompare */
610 0, /* tp_weaklistoffset */
611 0, /* tp_iter */
612 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000613 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000614 0, /* tp_members */
615 0, /* tp_getset */
616 0, /* tp_base */
617 0, /* tp_dict */
618 0, /* tp_descr_get */
619 0, /* tp_descr_set */
620 0, /* tp_dictoffset */
621 0, /* tp_init */
622 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000623 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000624 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000625};
626
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000627
628static PyTypeObject tee_type;
629
630static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000632{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000635 if (to->index >= LINKCELLS) {
636 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200637 if (link == NULL)
638 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000639 Py_DECREF(to->dataobj);
640 to->dataobj = (teedataobject *)link;
641 to->index = 0;
642 }
643 value = teedataobject_getitem(to->dataobj, to->index);
644 if (value == NULL)
645 return NULL;
646 to->index++;
647 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000648}
649
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000650static int
651tee_traverse(teeobject *to, visitproc visit, void *arg)
652{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000653 Py_VISIT((PyObject *)to->dataobj);
654 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000655}
656
Raymond Hettingerad983e72003-11-12 14:32:26 +0000657static PyObject *
658tee_copy(teeobject *to)
659{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000660 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 newto = PyObject_GC_New(teeobject, &tee_type);
663 if (newto == NULL)
664 return NULL;
665 Py_INCREF(to->dataobj);
666 newto->dataobj = to->dataobj;
667 newto->index = to->index;
668 newto->weakreflist = NULL;
669 PyObject_GC_Track(newto);
670 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000671}
672
673PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
674
675static PyObject *
676tee_fromiterable(PyObject *iterable)
677{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000678 teeobject *to;
679 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000680
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000681 it = PyObject_GetIter(iterable);
682 if (it == NULL)
683 return NULL;
684 if (PyObject_TypeCheck(it, &tee_type)) {
685 to = (teeobject *)tee_copy((teeobject *)it);
686 goto done;
687 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000688
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000689 to = PyObject_GC_New(teeobject, &tee_type);
690 if (to == NULL)
691 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000692 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 if (!to->dataobj) {
694 PyObject_GC_Del(to);
695 to = NULL;
696 goto done;
697 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000698
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000699 to->index = 0;
700 to->weakreflist = NULL;
701 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000702done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000703 Py_XDECREF(it);
704 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000705}
706
707static PyObject *
708tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000710 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000711
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000712 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000713 return NULL;
714 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000715}
716
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000717static int
718tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000719{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000720 if (to->weakreflist != NULL)
721 PyObject_ClearWeakRefs((PyObject *) to);
722 Py_CLEAR(to->dataobj);
723 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000724}
725
726static void
727tee_dealloc(teeobject *to)
728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000729 PyObject_GC_UnTrack(to);
730 tee_clear(to);
731 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000732}
733
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000734static PyObject *
735tee_reduce(teeobject *to)
736{
737 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
738}
739
740static PyObject *
741tee_setstate(teeobject *to, PyObject *state)
742{
743 teedataobject *tdo;
744 int index;
745 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
746 return NULL;
747 if (index < 0 || index > LINKCELLS) {
748 PyErr_SetString(PyExc_ValueError, "Index out of range");
749 return NULL;
750 }
751 Py_CLEAR(to->dataobj);
752 to->dataobj = tdo;
753 Py_INCREF(to->dataobj);
754 to->index = index;
755 Py_RETURN_NONE;
756}
757
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758PyDoc_STRVAR(teeobject_doc,
759"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000760
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000762 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000763 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
764 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000765 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000766};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000767
768static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000769 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000770 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000771 sizeof(teeobject), /* tp_basicsize */
772 0, /* tp_itemsize */
773 /* methods */
774 (destructor)tee_dealloc, /* tp_dealloc */
775 0, /* tp_print */
776 0, /* tp_getattr */
777 0, /* tp_setattr */
778 0, /* tp_reserved */
779 0, /* tp_repr */
780 0, /* tp_as_number */
781 0, /* tp_as_sequence */
782 0, /* tp_as_mapping */
783 0, /* tp_hash */
784 0, /* tp_call */
785 0, /* tp_str */
786 0, /* tp_getattro */
787 0, /* tp_setattro */
788 0, /* tp_as_buffer */
789 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
790 teeobject_doc, /* tp_doc */
791 (traverseproc)tee_traverse, /* tp_traverse */
792 (inquiry)tee_clear, /* tp_clear */
793 0, /* tp_richcompare */
794 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
795 PyObject_SelfIter, /* tp_iter */
796 (iternextfunc)tee_next, /* tp_iternext */
797 tee_methods, /* tp_methods */
798 0, /* tp_members */
799 0, /* tp_getset */
800 0, /* tp_base */
801 0, /* tp_dict */
802 0, /* tp_descr_get */
803 0, /* tp_descr_set */
804 0, /* tp_dictoffset */
805 0, /* tp_init */
806 0, /* tp_alloc */
807 tee_new, /* tp_new */
808 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000809};
810
Raymond Hettingerad983e72003-11-12 14:32:26 +0000811static PyObject *
812tee(PyObject *self, PyObject *args)
813{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000814 Py_ssize_t i, n=2;
815 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200816 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000817
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000818 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
819 return NULL;
820 if (n < 0) {
821 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
822 return NULL;
823 }
824 result = PyTuple_New(n);
825 if (result == NULL)
826 return NULL;
827 if (n == 0)
828 return result;
829 it = PyObject_GetIter(iterable);
830 if (it == NULL) {
831 Py_DECREF(result);
832 return NULL;
833 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200834 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000835 copyable = tee_fromiterable(it);
836 Py_DECREF(it);
837 if (copyable == NULL) {
838 Py_DECREF(result);
839 return NULL;
840 }
841 } else
842 copyable = it;
843 PyTuple_SET_ITEM(result, 0, copyable);
844 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200845
846 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000847 if (copyable == NULL) {
848 Py_DECREF(result);
849 return NULL;
850 }
851 PyTuple_SET_ITEM(result, i, copyable);
852 }
853 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000854}
855
856PyDoc_STRVAR(tee_doc,
857"tee(iterable, n=2) --> tuple of n independent iterators.");
858
859
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000860/* cycle object **********************************************************/
861
862typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 PyObject_HEAD
864 PyObject *it;
865 PyObject *saved;
866 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000867} cycleobject;
868
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000869static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000870
871static PyObject *
872cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000874 PyObject *it;
875 PyObject *iterable;
876 PyObject *saved;
877 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000879 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
880 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000881
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000882 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
883 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000885 /* Get iterator. */
886 it = PyObject_GetIter(iterable);
887 if (it == NULL)
888 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000889
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000890 saved = PyList_New(0);
891 if (saved == NULL) {
892 Py_DECREF(it);
893 return NULL;
894 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000895
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000896 /* create cycleobject structure */
897 lz = (cycleobject *)type->tp_alloc(type, 0);
898 if (lz == NULL) {
899 Py_DECREF(it);
900 Py_DECREF(saved);
901 return NULL;
902 }
903 lz->it = it;
904 lz->saved = saved;
905 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000907 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000908}
909
910static void
911cycle_dealloc(cycleobject *lz)
912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 PyObject_GC_UnTrack(lz);
914 Py_XDECREF(lz->saved);
915 Py_XDECREF(lz->it);
916 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000917}
918
919static int
920cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
921{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000922 Py_VISIT(lz->it);
923 Py_VISIT(lz->saved);
924 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000925}
926
927static PyObject *
928cycle_next(cycleobject *lz)
929{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000930 PyObject *item;
931 PyObject *it;
932 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000933
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000934 while (1) {
935 item = PyIter_Next(lz->it);
936 if (item != NULL) {
937 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
938 Py_DECREF(item);
939 return NULL;
940 }
941 return item;
942 }
943 if (PyErr_Occurred()) {
944 if (PyErr_ExceptionMatches(PyExc_StopIteration))
945 PyErr_Clear();
946 else
947 return NULL;
948 }
949 if (PyList_Size(lz->saved) == 0)
950 return NULL;
951 it = PyObject_GetIter(lz->saved);
952 if (it == NULL)
953 return NULL;
954 tmp = lz->it;
955 lz->it = it;
956 lz->firstpass = 1;
957 Py_DECREF(tmp);
958 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000959}
960
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000961static PyObject *
962cycle_reduce(cycleobject *lz)
963{
964 /* Create a new cycle with the iterator tuple, then set
965 * the saved state on it.
966 */
Serhiy Storchakad269b5e2013-01-25 13:38:56 +0200967 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000968 lz->it, lz->saved, lz->firstpass);
969 }
970
971static PyObject *
972cycle_setstate(cycleobject *lz, PyObject *state)
973{
974 PyObject *saved=NULL;
975 int firstpass;
976 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
977 return NULL;
978 Py_CLEAR(lz->saved);
979 lz->saved = saved;
980 Py_XINCREF(lz->saved);
981 lz->firstpass = firstpass != 0;
982 Py_RETURN_NONE;
983}
984
985static PyMethodDef cycle_methods[] = {
986 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
987 reduce_doc},
988 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
989 setstate_doc},
990 {NULL, NULL} /* sentinel */
991};
992
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000993PyDoc_STRVAR(cycle_doc,
994"cycle(iterable) --> cycle object\n\
995\n\
996Return elements from the iterable until it is exhausted.\n\
997Then repeat the sequence indefinitely.");
998
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000999static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 PyVarObject_HEAD_INIT(NULL, 0)
1001 "itertools.cycle", /* tp_name */
1002 sizeof(cycleobject), /* tp_basicsize */
1003 0, /* tp_itemsize */
1004 /* methods */
1005 (destructor)cycle_dealloc, /* tp_dealloc */
1006 0, /* tp_print */
1007 0, /* tp_getattr */
1008 0, /* tp_setattr */
1009 0, /* tp_reserved */
1010 0, /* tp_repr */
1011 0, /* tp_as_number */
1012 0, /* tp_as_sequence */
1013 0, /* tp_as_mapping */
1014 0, /* tp_hash */
1015 0, /* tp_call */
1016 0, /* tp_str */
1017 PyObject_GenericGetAttr, /* tp_getattro */
1018 0, /* tp_setattro */
1019 0, /* tp_as_buffer */
1020 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1021 Py_TPFLAGS_BASETYPE, /* tp_flags */
1022 cycle_doc, /* tp_doc */
1023 (traverseproc)cycle_traverse, /* tp_traverse */
1024 0, /* tp_clear */
1025 0, /* tp_richcompare */
1026 0, /* tp_weaklistoffset */
1027 PyObject_SelfIter, /* tp_iter */
1028 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001029 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001030 0, /* tp_members */
1031 0, /* tp_getset */
1032 0, /* tp_base */
1033 0, /* tp_dict */
1034 0, /* tp_descr_get */
1035 0, /* tp_descr_set */
1036 0, /* tp_dictoffset */
1037 0, /* tp_init */
1038 0, /* tp_alloc */
1039 cycle_new, /* tp_new */
1040 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001041};
1042
1043
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001044/* dropwhile object **********************************************************/
1045
1046typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001047 PyObject_HEAD
1048 PyObject *func;
1049 PyObject *it;
1050 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051} dropwhileobject;
1052
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001053static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001054
1055static PyObject *
1056dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1057{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001058 PyObject *func, *seq;
1059 PyObject *it;
1060 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001062 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1063 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1066 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001068 /* Get iterator. */
1069 it = PyObject_GetIter(seq);
1070 if (it == NULL)
1071 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001073 /* create dropwhileobject structure */
1074 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1075 if (lz == NULL) {
1076 Py_DECREF(it);
1077 return NULL;
1078 }
1079 Py_INCREF(func);
1080 lz->func = func;
1081 lz->it = it;
1082 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001084 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001085}
1086
1087static void
1088dropwhile_dealloc(dropwhileobject *lz)
1089{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001090 PyObject_GC_UnTrack(lz);
1091 Py_XDECREF(lz->func);
1092 Py_XDECREF(lz->it);
1093 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094}
1095
1096static int
1097dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1098{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001099 Py_VISIT(lz->it);
1100 Py_VISIT(lz->func);
1101 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001102}
1103
1104static PyObject *
1105dropwhile_next(dropwhileobject *lz)
1106{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 PyObject *item, *good;
1108 PyObject *it = lz->it;
1109 long ok;
1110 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001112 iternext = *Py_TYPE(it)->tp_iternext;
1113 for (;;) {
1114 item = iternext(it);
1115 if (item == NULL)
1116 return NULL;
1117 if (lz->start == 1)
1118 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1121 if (good == NULL) {
1122 Py_DECREF(item);
1123 return NULL;
1124 }
1125 ok = PyObject_IsTrue(good);
1126 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001127 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001128 lz->start = 1;
1129 return item;
1130 }
1131 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001132 if (ok < 0)
1133 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001134 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135}
1136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001137static PyObject *
1138dropwhile_reduce(dropwhileobject *lz)
1139{
1140 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1141 lz->func, lz->it, lz->start);
1142}
1143
1144static PyObject *
1145dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1146{
1147 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001148 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001149 return NULL;
1150 lz->start = start;
1151 Py_RETURN_NONE;
1152}
1153
1154static PyMethodDef dropwhile_methods[] = {
1155 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1156 reduce_doc},
1157 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1158 setstate_doc},
1159 {NULL, NULL} /* sentinel */
1160};
1161
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001162PyDoc_STRVAR(dropwhile_doc,
1163"dropwhile(predicate, iterable) --> dropwhile object\n\
1164\n\
1165Drop items from the iterable while predicate(item) is true.\n\
1166Afterwards, return every element until the iterable is exhausted.");
1167
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001168static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001169 PyVarObject_HEAD_INIT(NULL, 0)
1170 "itertools.dropwhile", /* tp_name */
1171 sizeof(dropwhileobject), /* tp_basicsize */
1172 0, /* tp_itemsize */
1173 /* methods */
1174 (destructor)dropwhile_dealloc, /* tp_dealloc */
1175 0, /* tp_print */
1176 0, /* tp_getattr */
1177 0, /* tp_setattr */
1178 0, /* tp_reserved */
1179 0, /* tp_repr */
1180 0, /* tp_as_number */
1181 0, /* tp_as_sequence */
1182 0, /* tp_as_mapping */
1183 0, /* tp_hash */
1184 0, /* tp_call */
1185 0, /* tp_str */
1186 PyObject_GenericGetAttr, /* tp_getattro */
1187 0, /* tp_setattro */
1188 0, /* tp_as_buffer */
1189 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1190 Py_TPFLAGS_BASETYPE, /* tp_flags */
1191 dropwhile_doc, /* tp_doc */
1192 (traverseproc)dropwhile_traverse, /* tp_traverse */
1193 0, /* tp_clear */
1194 0, /* tp_richcompare */
1195 0, /* tp_weaklistoffset */
1196 PyObject_SelfIter, /* tp_iter */
1197 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001198 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001199 0, /* tp_members */
1200 0, /* tp_getset */
1201 0, /* tp_base */
1202 0, /* tp_dict */
1203 0, /* tp_descr_get */
1204 0, /* tp_descr_set */
1205 0, /* tp_dictoffset */
1206 0, /* tp_init */
1207 0, /* tp_alloc */
1208 dropwhile_new, /* tp_new */
1209 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210};
1211
1212
1213/* takewhile object **********************************************************/
1214
1215typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001216 PyObject_HEAD
1217 PyObject *func;
1218 PyObject *it;
1219 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220} takewhileobject;
1221
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001222static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223
1224static PyObject *
1225takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1226{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001227 PyObject *func, *seq;
1228 PyObject *it;
1229 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1232 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1235 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001237 /* Get iterator. */
1238 it = PyObject_GetIter(seq);
1239 if (it == NULL)
1240 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001242 /* create takewhileobject structure */
1243 lz = (takewhileobject *)type->tp_alloc(type, 0);
1244 if (lz == NULL) {
1245 Py_DECREF(it);
1246 return NULL;
1247 }
1248 Py_INCREF(func);
1249 lz->func = func;
1250 lz->it = it;
1251 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001253 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254}
1255
1256static void
1257takewhile_dealloc(takewhileobject *lz)
1258{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001259 PyObject_GC_UnTrack(lz);
1260 Py_XDECREF(lz->func);
1261 Py_XDECREF(lz->it);
1262 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001263}
1264
1265static int
1266takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1267{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001268 Py_VISIT(lz->it);
1269 Py_VISIT(lz->func);
1270 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271}
1272
1273static PyObject *
1274takewhile_next(takewhileobject *lz)
1275{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001276 PyObject *item, *good;
1277 PyObject *it = lz->it;
1278 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001280 if (lz->stop == 1)
1281 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001282
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001283 item = (*Py_TYPE(it)->tp_iternext)(it);
1284 if (item == NULL)
1285 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1288 if (good == NULL) {
1289 Py_DECREF(item);
1290 return NULL;
1291 }
1292 ok = PyObject_IsTrue(good);
1293 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001294 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001295 return item;
1296 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001297 if (ok == 0)
1298 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001299 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300}
1301
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001302static PyObject *
1303takewhile_reduce(takewhileobject *lz)
1304{
1305 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1306 lz->func, lz->it, lz->stop);
1307}
1308
1309static PyObject *
1310takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1311{
1312 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001313 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001314 return NULL;
1315 lz->stop = stop;
1316 Py_RETURN_NONE;
1317}
1318
1319static PyMethodDef takewhile_reduce_methods[] = {
1320 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1321 reduce_doc},
1322 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1323 setstate_doc},
1324 {NULL, NULL} /* sentinel */
1325};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326PyDoc_STRVAR(takewhile_doc,
1327"takewhile(predicate, iterable) --> takewhile object\n\
1328\n\
1329Return successive entries from an iterable as long as the \n\
1330predicate evaluates to true for each entry.");
1331
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001332static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001333 PyVarObject_HEAD_INIT(NULL, 0)
1334 "itertools.takewhile", /* tp_name */
1335 sizeof(takewhileobject), /* tp_basicsize */
1336 0, /* tp_itemsize */
1337 /* methods */
1338 (destructor)takewhile_dealloc, /* tp_dealloc */
1339 0, /* tp_print */
1340 0, /* tp_getattr */
1341 0, /* tp_setattr */
1342 0, /* tp_reserved */
1343 0, /* tp_repr */
1344 0, /* tp_as_number */
1345 0, /* tp_as_sequence */
1346 0, /* tp_as_mapping */
1347 0, /* tp_hash */
1348 0, /* tp_call */
1349 0, /* tp_str */
1350 PyObject_GenericGetAttr, /* tp_getattro */
1351 0, /* tp_setattro */
1352 0, /* tp_as_buffer */
1353 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1354 Py_TPFLAGS_BASETYPE, /* tp_flags */
1355 takewhile_doc, /* tp_doc */
1356 (traverseproc)takewhile_traverse, /* tp_traverse */
1357 0, /* tp_clear */
1358 0, /* tp_richcompare */
1359 0, /* tp_weaklistoffset */
1360 PyObject_SelfIter, /* tp_iter */
1361 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001362 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 0, /* tp_members */
1364 0, /* tp_getset */
1365 0, /* tp_base */
1366 0, /* tp_dict */
1367 0, /* tp_descr_get */
1368 0, /* tp_descr_set */
1369 0, /* tp_dictoffset */
1370 0, /* tp_init */
1371 0, /* tp_alloc */
1372 takewhile_new, /* tp_new */
1373 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374};
1375
1376
1377/* islice object ************************************************************/
1378
1379typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001380 PyObject_HEAD
1381 PyObject *it;
1382 Py_ssize_t next;
1383 Py_ssize_t stop;
1384 Py_ssize_t step;
1385 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386} isliceobject;
1387
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001388static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001389
1390static PyObject *
1391islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1392{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001393 PyObject *seq;
1394 Py_ssize_t start=0, stop=-1, step=1;
1395 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1396 Py_ssize_t numargs;
1397 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001399 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1400 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001401
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001402 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1403 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001405 numargs = PyTuple_Size(args);
1406 if (numargs == 2) {
1407 if (a1 != Py_None) {
1408 stop = PyLong_AsSsize_t(a1);
1409 if (stop == -1) {
1410 if (PyErr_Occurred())
1411 PyErr_Clear();
1412 PyErr_SetString(PyExc_ValueError,
1413 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1414 return NULL;
1415 }
1416 }
1417 } else {
1418 if (a1 != Py_None)
1419 start = PyLong_AsSsize_t(a1);
1420 if (start == -1 && PyErr_Occurred())
1421 PyErr_Clear();
1422 if (a2 != Py_None) {
1423 stop = PyLong_AsSsize_t(a2);
1424 if (stop == -1) {
1425 if (PyErr_Occurred())
1426 PyErr_Clear();
1427 PyErr_SetString(PyExc_ValueError,
1428 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1429 return NULL;
1430 }
1431 }
1432 }
1433 if (start<0 || stop<-1) {
1434 PyErr_SetString(PyExc_ValueError,
1435 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1436 return NULL;
1437 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001439 if (a3 != NULL) {
1440 if (a3 != Py_None)
1441 step = PyLong_AsSsize_t(a3);
1442 if (step == -1 && PyErr_Occurred())
1443 PyErr_Clear();
1444 }
1445 if (step<1) {
1446 PyErr_SetString(PyExc_ValueError,
1447 "Step for islice() must be a positive integer or None.");
1448 return NULL;
1449 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001451 /* Get iterator. */
1452 it = PyObject_GetIter(seq);
1453 if (it == NULL)
1454 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 /* create isliceobject structure */
1457 lz = (isliceobject *)type->tp_alloc(type, 0);
1458 if (lz == NULL) {
1459 Py_DECREF(it);
1460 return NULL;
1461 }
1462 lz->it = it;
1463 lz->next = start;
1464 lz->stop = stop;
1465 lz->step = step;
1466 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001468 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469}
1470
1471static void
1472islice_dealloc(isliceobject *lz)
1473{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001474 PyObject_GC_UnTrack(lz);
1475 Py_XDECREF(lz->it);
1476 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477}
1478
1479static int
1480islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 Py_VISIT(lz->it);
1483 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static PyObject *
1487islice_next(isliceobject *lz)
1488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 PyObject *item;
1490 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001491 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 Py_ssize_t oldnext;
1493 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001495 if (it == NULL)
1496 return NULL;
1497
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 iternext = *Py_TYPE(it)->tp_iternext;
1499 while (lz->cnt < lz->next) {
1500 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 Py_DECREF(item);
1504 lz->cnt++;
1505 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001506 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001507 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 item = iternext(it);
1509 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001510 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001511 lz->cnt++;
1512 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001513 /* The (size_t) cast below avoids the danger of undefined
1514 behaviour from signed integer overflow. */
1515 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001516 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1517 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001518 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001519
1520empty:
1521 Py_CLEAR(lz->it);
1522 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523}
1524
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001525static PyObject *
1526islice_reduce(isliceobject *lz)
1527{
1528 /* When unpickled, generate a new object with the same bounds,
1529 * then 'setstate' with the next and count
1530 */
1531 PyObject *stop;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001532 if (lz->it == NULL) {
1533 PyObject *empty_list;
1534 PyObject *empty_it;
1535 empty_list = PyList_New(0);
1536 if (empty_list == NULL)
1537 return NULL;
1538 empty_it = PyObject_GetIter(empty_list);
1539 Py_DECREF(empty_list);
1540 if (empty_it == NULL)
1541 return NULL;
1542 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1543 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001544 if (lz->stop == -1) {
1545 stop = Py_None;
1546 Py_INCREF(stop);
1547 } else {
1548 stop = PyLong_FromSsize_t(lz->stop);
1549 if (stop == NULL)
1550 return NULL;
1551 }
1552 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1553 lz->it, lz->next, stop, lz->step,
1554 lz->cnt);
1555}
1556
1557static PyObject *
1558islice_setstate(isliceobject *lz, PyObject *state)
1559{
1560 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1561 if (cnt == -1 && PyErr_Occurred())
1562 return NULL;
1563 lz->cnt = cnt;
1564 Py_RETURN_NONE;
1565}
1566
1567static PyMethodDef islice_methods[] = {
1568 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1569 reduce_doc},
1570 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1571 setstate_doc},
1572 {NULL, NULL} /* sentinel */
1573};
1574
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001575PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001576"islice(iterable, stop) --> islice object\n\
1577islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578\n\
1579Return an iterator whose next() method returns selected values from an\n\
1580iterable. If start is specified, will skip all preceding elements;\n\
1581otherwise, start defaults to zero. Step defaults to one. If\n\
1582specified as another value, step determines how many values are \n\
1583skipped between successive calls. Works like a slice() on a list\n\
1584but returns an iterator.");
1585
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001586static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001587 PyVarObject_HEAD_INIT(NULL, 0)
1588 "itertools.islice", /* tp_name */
1589 sizeof(isliceobject), /* tp_basicsize */
1590 0, /* tp_itemsize */
1591 /* methods */
1592 (destructor)islice_dealloc, /* tp_dealloc */
1593 0, /* tp_print */
1594 0, /* tp_getattr */
1595 0, /* tp_setattr */
1596 0, /* tp_reserved */
1597 0, /* tp_repr */
1598 0, /* tp_as_number */
1599 0, /* tp_as_sequence */
1600 0, /* tp_as_mapping */
1601 0, /* tp_hash */
1602 0, /* tp_call */
1603 0, /* tp_str */
1604 PyObject_GenericGetAttr, /* tp_getattro */
1605 0, /* tp_setattro */
1606 0, /* tp_as_buffer */
1607 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1608 Py_TPFLAGS_BASETYPE, /* tp_flags */
1609 islice_doc, /* tp_doc */
1610 (traverseproc)islice_traverse, /* tp_traverse */
1611 0, /* tp_clear */
1612 0, /* tp_richcompare */
1613 0, /* tp_weaklistoffset */
1614 PyObject_SelfIter, /* tp_iter */
1615 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001616 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001617 0, /* tp_members */
1618 0, /* tp_getset */
1619 0, /* tp_base */
1620 0, /* tp_dict */
1621 0, /* tp_descr_get */
1622 0, /* tp_descr_set */
1623 0, /* tp_dictoffset */
1624 0, /* tp_init */
1625 0, /* tp_alloc */
1626 islice_new, /* tp_new */
1627 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001628};
1629
1630
1631/* starmap object ************************************************************/
1632
1633typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001634 PyObject_HEAD
1635 PyObject *func;
1636 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637} starmapobject;
1638
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001639static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640
1641static PyObject *
1642starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1643{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001644 PyObject *func, *seq;
1645 PyObject *it;
1646 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001648 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1649 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001651 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1652 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001653
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001654 /* Get iterator. */
1655 it = PyObject_GetIter(seq);
1656 if (it == NULL)
1657 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001659 /* create starmapobject structure */
1660 lz = (starmapobject *)type->tp_alloc(type, 0);
1661 if (lz == NULL) {
1662 Py_DECREF(it);
1663 return NULL;
1664 }
1665 Py_INCREF(func);
1666 lz->func = func;
1667 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001669 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670}
1671
1672static void
1673starmap_dealloc(starmapobject *lz)
1674{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001675 PyObject_GC_UnTrack(lz);
1676 Py_XDECREF(lz->func);
1677 Py_XDECREF(lz->it);
1678 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679}
1680
1681static int
1682starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1683{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001684 Py_VISIT(lz->it);
1685 Py_VISIT(lz->func);
1686 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687}
1688
1689static PyObject *
1690starmap_next(starmapobject *lz)
1691{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001692 PyObject *args;
1693 PyObject *result;
1694 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001696 args = (*Py_TYPE(it)->tp_iternext)(it);
1697 if (args == NULL)
1698 return NULL;
1699 if (!PyTuple_CheckExact(args)) {
1700 PyObject *newargs = PySequence_Tuple(args);
1701 Py_DECREF(args);
1702 if (newargs == NULL)
1703 return NULL;
1704 args = newargs;
1705 }
1706 result = PyObject_Call(lz->func, args, NULL);
1707 Py_DECREF(args);
1708 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709}
1710
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001711static PyObject *
1712starmap_reduce(starmapobject *lz)
1713{
1714 /* Just pickle the iterator */
1715 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1716}
1717
1718static PyMethodDef starmap_methods[] = {
1719 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1720 reduce_doc},
1721 {NULL, NULL} /* sentinel */
1722};
1723
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724PyDoc_STRVAR(starmap_doc,
1725"starmap(function, sequence) --> starmap object\n\
1726\n\
1727Return an iterator whose values are returned from the function evaluated\n\
1728with a argument tuple taken from the given sequence.");
1729
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001730static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001731 PyVarObject_HEAD_INIT(NULL, 0)
1732 "itertools.starmap", /* tp_name */
1733 sizeof(starmapobject), /* tp_basicsize */
1734 0, /* tp_itemsize */
1735 /* methods */
1736 (destructor)starmap_dealloc, /* tp_dealloc */
1737 0, /* tp_print */
1738 0, /* tp_getattr */
1739 0, /* tp_setattr */
1740 0, /* tp_reserved */
1741 0, /* tp_repr */
1742 0, /* tp_as_number */
1743 0, /* tp_as_sequence */
1744 0, /* tp_as_mapping */
1745 0, /* tp_hash */
1746 0, /* tp_call */
1747 0, /* tp_str */
1748 PyObject_GenericGetAttr, /* tp_getattro */
1749 0, /* tp_setattro */
1750 0, /* tp_as_buffer */
1751 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1752 Py_TPFLAGS_BASETYPE, /* tp_flags */
1753 starmap_doc, /* tp_doc */
1754 (traverseproc)starmap_traverse, /* tp_traverse */
1755 0, /* tp_clear */
1756 0, /* tp_richcompare */
1757 0, /* tp_weaklistoffset */
1758 PyObject_SelfIter, /* tp_iter */
1759 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001760 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001761 0, /* tp_members */
1762 0, /* tp_getset */
1763 0, /* tp_base */
1764 0, /* tp_dict */
1765 0, /* tp_descr_get */
1766 0, /* tp_descr_set */
1767 0, /* tp_dictoffset */
1768 0, /* tp_init */
1769 0, /* tp_alloc */
1770 starmap_new, /* tp_new */
1771 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001772};
1773
1774
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001775/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001776
1777typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001778 PyObject_HEAD
1779 PyObject *source; /* Iterator over input iterables */
1780 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001781} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001782
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001783static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001786chain_new_internal(PyTypeObject *type, PyObject *source)
1787{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001788 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001789
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001790 lz = (chainobject *)type->tp_alloc(type, 0);
1791 if (lz == NULL) {
1792 Py_DECREF(source);
1793 return NULL;
1794 }
1795
1796 lz->source = source;
1797 lz->active = NULL;
1798 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001799}
1800
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001802chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001804 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001806 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1807 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 source = PyObject_GetIter(args);
1810 if (source == NULL)
1811 return NULL;
1812
1813 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001814}
1815
1816static PyObject *
1817chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1818{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001819 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 source = PyObject_GetIter(arg);
1822 if (source == NULL)
1823 return NULL;
1824
1825 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826}
1827
1828static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001829chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001831 PyObject_GC_UnTrack(lz);
1832 Py_XDECREF(lz->active);
1833 Py_XDECREF(lz->source);
1834 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001835}
1836
Raymond Hettinger2012f172003-02-07 05:32:58 +00001837static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001838chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001840 Py_VISIT(lz->source);
1841 Py_VISIT(lz->active);
1842 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001843}
1844
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001845static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001846chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001848 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001850 if (lz->source == NULL)
1851 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001852
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001853 if (lz->active == NULL) {
1854 PyObject *iterable = PyIter_Next(lz->source);
1855 if (iterable == NULL) {
1856 Py_CLEAR(lz->source);
1857 return NULL; /* no more input sources */
1858 }
1859 lz->active = PyObject_GetIter(iterable);
1860 Py_DECREF(iterable);
1861 if (lz->active == NULL) {
1862 Py_CLEAR(lz->source);
1863 return NULL; /* input not iterable */
1864 }
1865 }
1866 item = PyIter_Next(lz->active);
1867 if (item != NULL)
1868 return item;
1869 if (PyErr_Occurred()) {
1870 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1871 PyErr_Clear();
1872 else
1873 return NULL; /* input raised an exception */
1874 }
1875 Py_CLEAR(lz->active);
1876 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001877}
1878
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001879static PyObject *
1880chain_reduce(chainobject *lz)
1881{
1882 if (lz->source) {
1883 /* we can't pickle function objects (itertools.from_iterable) so
1884 * we must use setstate to replace the iterable. One day we
1885 * will fix pickling of functions
1886 */
1887 if (lz->active) {
1888 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1889 } else {
1890 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1891 }
1892 } else {
1893 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1894 }
1895 return NULL;
1896}
1897
1898static PyObject *
1899chain_setstate(chainobject *lz, PyObject *state)
1900{
1901 PyObject *source, *active=NULL;
1902 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1903 return NULL;
1904
1905 Py_CLEAR(lz->source);
1906 lz->source = source;
1907 Py_INCREF(lz->source);
1908 Py_CLEAR(lz->active);
1909 lz->active = active;
1910 Py_XINCREF(lz->active);
1911 Py_RETURN_NONE;
1912}
1913
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001914PyDoc_STRVAR(chain_doc,
1915"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001916\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001917Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001918first iterable until it is exhausted, then elements from the next\n\
1919iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001920
Christian Heimesf16baeb2008-02-29 14:57:44 +00001921PyDoc_STRVAR(chain_from_iterable_doc,
1922"chain.from_iterable(iterable) --> chain object\n\
1923\n\
1924Alternate chain() contructor taking a single iterable argument\n\
1925that evaluates lazily.");
1926
1927static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001928 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1929 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001930 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1931 reduce_doc},
1932 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1933 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001934 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001935};
1936
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001937static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001938 PyVarObject_HEAD_INIT(NULL, 0)
1939 "itertools.chain", /* tp_name */
1940 sizeof(chainobject), /* tp_basicsize */
1941 0, /* tp_itemsize */
1942 /* methods */
1943 (destructor)chain_dealloc, /* tp_dealloc */
1944 0, /* tp_print */
1945 0, /* tp_getattr */
1946 0, /* tp_setattr */
1947 0, /* tp_reserved */
1948 0, /* tp_repr */
1949 0, /* tp_as_number */
1950 0, /* tp_as_sequence */
1951 0, /* tp_as_mapping */
1952 0, /* tp_hash */
1953 0, /* tp_call */
1954 0, /* tp_str */
1955 PyObject_GenericGetAttr, /* tp_getattro */
1956 0, /* tp_setattro */
1957 0, /* tp_as_buffer */
1958 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1959 Py_TPFLAGS_BASETYPE, /* tp_flags */
1960 chain_doc, /* tp_doc */
1961 (traverseproc)chain_traverse, /* tp_traverse */
1962 0, /* tp_clear */
1963 0, /* tp_richcompare */
1964 0, /* tp_weaklistoffset */
1965 PyObject_SelfIter, /* tp_iter */
1966 (iternextfunc)chain_next, /* tp_iternext */
1967 chain_methods, /* tp_methods */
1968 0, /* tp_members */
1969 0, /* tp_getset */
1970 0, /* tp_base */
1971 0, /* tp_dict */
1972 0, /* tp_descr_get */
1973 0, /* tp_descr_set */
1974 0, /* tp_dictoffset */
1975 0, /* tp_init */
1976 0, /* tp_alloc */
1977 chain_new, /* tp_new */
1978 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001979};
1980
1981
Christian Heimesc3f30c42008-02-22 16:37:40 +00001982/* product object ************************************************************/
1983
1984typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001985 PyObject_HEAD
1986 PyObject *pools; /* tuple of pool tuples */
1987 Py_ssize_t *indices; /* one index per pool */
1988 PyObject *result; /* most recently returned result tuple */
1989 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001990} productobject;
1991
1992static PyTypeObject product_type;
1993
1994static PyObject *
1995product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1996{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001997 productobject *lz;
1998 Py_ssize_t nargs, npools, repeat=1;
1999 PyObject *pools = NULL;
2000 Py_ssize_t *indices = NULL;
2001 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002002
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002003 if (kwds != NULL) {
2004 char *kwlist[] = {"repeat", 0};
2005 PyObject *tmpargs = PyTuple_New(0);
2006 if (tmpargs == NULL)
2007 return NULL;
2008 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
2009 Py_DECREF(tmpargs);
2010 return NULL;
2011 }
2012 Py_DECREF(tmpargs);
2013 if (repeat < 0) {
2014 PyErr_SetString(PyExc_ValueError,
2015 "repeat argument cannot be negative");
2016 return NULL;
2017 }
2018 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002019
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002020 assert(PyTuple_Check(args));
2021 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
2022 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002023
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002024 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
2025 if (indices == NULL) {
2026 PyErr_NoMemory();
2027 goto error;
2028 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002029
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002030 pools = PyTuple_New(npools);
2031 if (pools == NULL)
2032 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002034 for (i=0; i < nargs ; ++i) {
2035 PyObject *item = PyTuple_GET_ITEM(args, i);
2036 PyObject *pool = PySequence_Tuple(item);
2037 if (pool == NULL)
2038 goto error;
2039 PyTuple_SET_ITEM(pools, i, pool);
2040 indices[i] = 0;
2041 }
2042 for ( ; i < npools; ++i) {
2043 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2044 Py_INCREF(pool);
2045 PyTuple_SET_ITEM(pools, i, pool);
2046 indices[i] = 0;
2047 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002048
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002049 /* create productobject structure */
2050 lz = (productobject *)type->tp_alloc(type, 0);
2051 if (lz == NULL)
2052 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002054 lz->pools = pools;
2055 lz->indices = indices;
2056 lz->result = NULL;
2057 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002058
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002059 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002060
2061error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002062 if (indices != NULL)
2063 PyMem_Free(indices);
2064 Py_XDECREF(pools);
2065 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002066}
2067
2068static void
2069product_dealloc(productobject *lz)
2070{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002071 PyObject_GC_UnTrack(lz);
2072 Py_XDECREF(lz->pools);
2073 Py_XDECREF(lz->result);
2074 if (lz->indices != NULL)
2075 PyMem_Free(lz->indices);
2076 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002077}
2078
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002079static PyObject *
2080product_sizeof(productobject *lz, void *unused)
2081{
2082 Py_ssize_t res;
2083
2084 res = sizeof(productobject);
2085 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2086 return PyLong_FromSsize_t(res);
2087}
2088
2089PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2090
Christian Heimesc3f30c42008-02-22 16:37:40 +00002091static int
2092product_traverse(productobject *lz, visitproc visit, void *arg)
2093{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002094 Py_VISIT(lz->pools);
2095 Py_VISIT(lz->result);
2096 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002097}
2098
2099static PyObject *
2100product_next(productobject *lz)
2101{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002102 PyObject *pool;
2103 PyObject *elem;
2104 PyObject *oldelem;
2105 PyObject *pools = lz->pools;
2106 PyObject *result = lz->result;
2107 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2108 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002109
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002110 if (lz->stopped)
2111 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002112
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002113 if (result == NULL) {
2114 /* On the first pass, return an initial tuple filled with the
2115 first element from each pool. */
2116 result = PyTuple_New(npools);
2117 if (result == NULL)
2118 goto empty;
2119 lz->result = result;
2120 for (i=0; i < npools; i++) {
2121 pool = PyTuple_GET_ITEM(pools, i);
2122 if (PyTuple_GET_SIZE(pool) == 0)
2123 goto empty;
2124 elem = PyTuple_GET_ITEM(pool, 0);
2125 Py_INCREF(elem);
2126 PyTuple_SET_ITEM(result, i, elem);
2127 }
2128 } else {
2129 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002130
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002131 /* Copy the previous result tuple or re-use it if available */
2132 if (Py_REFCNT(result) > 1) {
2133 PyObject *old_result = result;
2134 result = PyTuple_New(npools);
2135 if (result == NULL)
2136 goto empty;
2137 lz->result = result;
2138 for (i=0; i < npools; i++) {
2139 elem = PyTuple_GET_ITEM(old_result, i);
2140 Py_INCREF(elem);
2141 PyTuple_SET_ITEM(result, i, elem);
2142 }
2143 Py_DECREF(old_result);
2144 }
2145 /* Now, we've got the only copy so we can update it in-place */
2146 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002147
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002148 /* Update the pool indices right-to-left. Only advance to the
2149 next pool when the previous one rolls-over */
2150 for (i=npools-1 ; i >= 0 ; i--) {
2151 pool = PyTuple_GET_ITEM(pools, i);
2152 indices[i]++;
2153 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2154 /* Roll-over and advance to next pool */
2155 indices[i] = 0;
2156 elem = PyTuple_GET_ITEM(pool, 0);
2157 Py_INCREF(elem);
2158 oldelem = PyTuple_GET_ITEM(result, i);
2159 PyTuple_SET_ITEM(result, i, elem);
2160 Py_DECREF(oldelem);
2161 } else {
2162 /* No rollover. Just increment and stop here. */
2163 elem = PyTuple_GET_ITEM(pool, indices[i]);
2164 Py_INCREF(elem);
2165 oldelem = PyTuple_GET_ITEM(result, i);
2166 PyTuple_SET_ITEM(result, i, elem);
2167 Py_DECREF(oldelem);
2168 break;
2169 }
2170 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002171
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002172 /* If i is negative, then the indices have all rolled-over
2173 and we're done. */
2174 if (i < 0)
2175 goto empty;
2176 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002178 Py_INCREF(result);
2179 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002180
2181empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002182 lz->stopped = 1;
2183 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002184}
2185
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002186static PyObject *
2187product_reduce(productobject *lz)
2188{
2189 if (lz->stopped) {
2190 return Py_BuildValue("O(())", Py_TYPE(lz));
2191 } else if (lz->result == NULL) {
2192 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2193 } else {
2194 PyObject *indices;
2195 Py_ssize_t n, i;
2196
2197 /* we must pickle the indices use them for setstate, and
2198 * additionally indicate that the iterator has started
2199 */
2200 n = PyTuple_GET_SIZE(lz->pools);
2201 indices = PyTuple_New(n);
2202 if (indices == NULL)
2203 return NULL;
2204 for (i=0; i<n; i++){
2205 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2206 if (!index) {
2207 Py_DECREF(indices);
2208 return NULL;
2209 }
2210 PyTuple_SET_ITEM(indices, i, index);
2211 }
2212 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2213 }
2214}
2215
2216static PyObject *
2217product_setstate(productobject *lz, PyObject *state)
2218{
2219 PyObject *result;
2220 Py_ssize_t n, i;
2221
2222 n = PyTuple_GET_SIZE(lz->pools);
2223 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2224 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2225 return NULL;
2226 }
2227 for (i=0; i<n; i++)
2228 {
2229 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2230 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2231 if (index < 0 && PyErr_Occurred())
2232 return NULL; /* not an integer */
2233 /* clamp the index */
2234 if (index < 0)
2235 index = 0;
2236 else if (index > n-1)
2237 index = n-1;
2238 lz->indices[i] = index;
2239 }
2240
2241 result = PyTuple_New(n);
2242 if (!result)
2243 return NULL;
2244 for (i=0; i<n; i++) {
2245 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2246 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2247 Py_INCREF(element);
2248 PyTuple_SET_ITEM(result, i, element);
2249 }
2250 Py_CLEAR(lz->result);
2251 lz->result = result;
2252 Py_RETURN_NONE;
2253}
2254
2255static PyMethodDef product_methods[] = {
2256 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2257 reduce_doc},
2258 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2259 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002260 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2261 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002262 {NULL, NULL} /* sentinel */
2263};
2264
Christian Heimesc3f30c42008-02-22 16:37:40 +00002265PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002266"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002267\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002268Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002269For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2270The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2271cycle in a manner similar to an odometer (with the rightmost element changing\n\
2272on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002273To compute the product of an iterable with itself, specify the number\n\
2274of repetitions with the optional repeat keyword argument. For example,\n\
2275product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002276product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2277product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2278
2279static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002280 PyVarObject_HEAD_INIT(NULL, 0)
2281 "itertools.product", /* tp_name */
2282 sizeof(productobject), /* tp_basicsize */
2283 0, /* tp_itemsize */
2284 /* methods */
2285 (destructor)product_dealloc, /* tp_dealloc */
2286 0, /* tp_print */
2287 0, /* tp_getattr */
2288 0, /* tp_setattr */
2289 0, /* tp_reserved */
2290 0, /* tp_repr */
2291 0, /* tp_as_number */
2292 0, /* tp_as_sequence */
2293 0, /* tp_as_mapping */
2294 0, /* tp_hash */
2295 0, /* tp_call */
2296 0, /* tp_str */
2297 PyObject_GenericGetAttr, /* tp_getattro */
2298 0, /* tp_setattro */
2299 0, /* tp_as_buffer */
2300 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2301 Py_TPFLAGS_BASETYPE, /* tp_flags */
2302 product_doc, /* tp_doc */
2303 (traverseproc)product_traverse, /* tp_traverse */
2304 0, /* tp_clear */
2305 0, /* tp_richcompare */
2306 0, /* tp_weaklistoffset */
2307 PyObject_SelfIter, /* tp_iter */
2308 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002309 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002310 0, /* tp_members */
2311 0, /* tp_getset */
2312 0, /* tp_base */
2313 0, /* tp_dict */
2314 0, /* tp_descr_get */
2315 0, /* tp_descr_set */
2316 0, /* tp_dictoffset */
2317 0, /* tp_init */
2318 0, /* tp_alloc */
2319 product_new, /* tp_new */
2320 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002321};
2322
2323
Christian Heimes380f7f22008-02-28 11:19:05 +00002324/* combinations object ************************************************************/
2325
2326typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002327 PyObject_HEAD
2328 PyObject *pool; /* input converted to a tuple */
2329 Py_ssize_t *indices; /* one index per result element */
2330 PyObject *result; /* most recently returned result tuple */
2331 Py_ssize_t r; /* size of result tuple */
2332 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002333} combinationsobject;
2334
2335static PyTypeObject combinations_type;
2336
2337static PyObject *
2338combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2339{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002340 combinationsobject *co;
2341 Py_ssize_t n;
2342 Py_ssize_t r;
2343 PyObject *pool = NULL;
2344 PyObject *iterable = NULL;
2345 Py_ssize_t *indices = NULL;
2346 Py_ssize_t i;
2347 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002348
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002349 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2350 &iterable, &r))
2351 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002352
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002353 pool = PySequence_Tuple(iterable);
2354 if (pool == NULL)
2355 goto error;
2356 n = PyTuple_GET_SIZE(pool);
2357 if (r < 0) {
2358 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2359 goto error;
2360 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002361
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002362 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2363 if (indices == NULL) {
2364 PyErr_NoMemory();
2365 goto error;
2366 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002368 for (i=0 ; i<r ; i++)
2369 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002370
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002371 /* create combinationsobject structure */
2372 co = (combinationsobject *)type->tp_alloc(type, 0);
2373 if (co == NULL)
2374 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002375
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 co->pool = pool;
2377 co->indices = indices;
2378 co->result = NULL;
2379 co->r = r;
2380 co->stopped = r > n ? 1 : 0;
2381
2382 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002383
2384error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002385 if (indices != NULL)
2386 PyMem_Free(indices);
2387 Py_XDECREF(pool);
2388 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002389}
2390
2391static void
2392combinations_dealloc(combinationsobject *co)
2393{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002394 PyObject_GC_UnTrack(co);
2395 Py_XDECREF(co->pool);
2396 Py_XDECREF(co->result);
2397 if (co->indices != NULL)
2398 PyMem_Free(co->indices);
2399 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002400}
2401
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002402static PyObject *
2403combinations_sizeof(combinationsobject *co, void *unused)
2404{
2405 Py_ssize_t res;
2406
2407 res = sizeof(combinationsobject);
2408 res += co->r * sizeof(Py_ssize_t);
2409 return PyLong_FromSsize_t(res);
2410}
2411
Christian Heimes380f7f22008-02-28 11:19:05 +00002412static int
2413combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2414{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002415 Py_VISIT(co->pool);
2416 Py_VISIT(co->result);
2417 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002418}
2419
2420static PyObject *
2421combinations_next(combinationsobject *co)
2422{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002423 PyObject *elem;
2424 PyObject *oldelem;
2425 PyObject *pool = co->pool;
2426 Py_ssize_t *indices = co->indices;
2427 PyObject *result = co->result;
2428 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2429 Py_ssize_t r = co->r;
2430 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002431
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002432 if (co->stopped)
2433 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002434
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 if (result == NULL) {
2436 /* On the first pass, initialize result tuple using the indices */
2437 result = PyTuple_New(r);
2438 if (result == NULL)
2439 goto empty;
2440 co->result = result;
2441 for (i=0; i<r ; i++) {
2442 index = indices[i];
2443 elem = PyTuple_GET_ITEM(pool, index);
2444 Py_INCREF(elem);
2445 PyTuple_SET_ITEM(result, i, elem);
2446 }
2447 } else {
2448 /* Copy the previous result tuple or re-use it if available */
2449 if (Py_REFCNT(result) > 1) {
2450 PyObject *old_result = result;
2451 result = PyTuple_New(r);
2452 if (result == NULL)
2453 goto empty;
2454 co->result = result;
2455 for (i=0; i<r ; i++) {
2456 elem = PyTuple_GET_ITEM(old_result, i);
2457 Py_INCREF(elem);
2458 PyTuple_SET_ITEM(result, i, elem);
2459 }
2460 Py_DECREF(old_result);
2461 }
2462 /* Now, we've got the only copy so we can update it in-place
2463 * CPython's empty tuple is a singleton and cached in
2464 * PyTuple's freelist.
2465 */
2466 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002468 /* Scan indices right-to-left until finding one that is not
2469 at its maximum (i + n - r). */
2470 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2471 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002472
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002473 /* If i is negative, then the indices are all at
2474 their maximum value and we're done. */
2475 if (i < 0)
2476 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002477
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002478 /* Increment the current index which we know is not at its
2479 maximum. Then move back to the right setting each index
2480 to its lowest possible value (one higher than the index
2481 to its left -- this maintains the sort order invariant). */
2482 indices[i]++;
2483 for (j=i+1 ; j<r ; j++)
2484 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002486 /* Update the result tuple for the new indices
2487 starting with i, the leftmost index that changed */
2488 for ( ; i<r ; i++) {
2489 index = indices[i];
2490 elem = PyTuple_GET_ITEM(pool, index);
2491 Py_INCREF(elem);
2492 oldelem = PyTuple_GET_ITEM(result, i);
2493 PyTuple_SET_ITEM(result, i, elem);
2494 Py_DECREF(oldelem);
2495 }
2496 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002497
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002498 Py_INCREF(result);
2499 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002500
2501empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002502 co->stopped = 1;
2503 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002504}
2505
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002506static PyObject *
2507combinations_reduce(combinationsobject *lz)
2508{
2509 if (lz->result == NULL) {
2510 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2511 } else if (lz->stopped) {
2512 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2513 } else {
2514 PyObject *indices;
2515 Py_ssize_t i;
2516
2517 /* we must pickle the indices and use them for setstate */
2518 indices = PyTuple_New(lz->r);
2519 if (!indices)
2520 return NULL;
2521 for (i=0; i<lz->r; i++)
2522 {
2523 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2524 if (!index) {
2525 Py_DECREF(indices);
2526 return NULL;
2527 }
2528 PyTuple_SET_ITEM(indices, i, index);
2529 }
2530
2531 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2532 }
2533}
2534
2535static PyObject *
2536combinations_setstate(combinationsobject *lz, PyObject *state)
2537{
2538 PyObject *result;
2539 Py_ssize_t i;
2540 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2541
2542 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2543 {
2544 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2545 return NULL;
2546 }
2547
2548 for (i=0; i<lz->r; i++)
2549 {
2550 Py_ssize_t max;
2551 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2552 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2553 if (index == -1 && PyErr_Occurred())
2554 return NULL; /* not an integer */
2555 max = i + n - lz->r;
2556 /* clamp the index (beware of negative max) */
2557 if (index > max)
2558 index = max;
2559 if (index < 0)
2560 index = 0;
2561 lz->indices[i] = index;
2562 }
2563
2564 result = PyTuple_New(lz->r);
2565 if (result == NULL)
2566 return NULL;
2567 for (i=0; i<lz->r; i++) {
2568 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2569 Py_INCREF(element);
2570 PyTuple_SET_ITEM(result, i, element);
2571 }
2572
2573 Py_CLEAR(lz->result);
2574 lz->result = result;
2575 Py_RETURN_NONE;
2576}
2577
2578static PyMethodDef combinations_methods[] = {
2579 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2580 reduce_doc},
2581 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2582 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002583 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2584 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002585 {NULL, NULL} /* sentinel */
2586};
2587
Christian Heimes380f7f22008-02-28 11:19:05 +00002588PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002589"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002590\n\
2591Return successive r-length combinations of elements in the iterable.\n\n\
2592combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2593
2594static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002595 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002596 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002597 sizeof(combinationsobject), /* tp_basicsize */
2598 0, /* tp_itemsize */
2599 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002600 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002601 0, /* tp_print */
2602 0, /* tp_getattr */
2603 0, /* tp_setattr */
2604 0, /* tp_reserved */
2605 0, /* tp_repr */
2606 0, /* tp_as_number */
2607 0, /* tp_as_sequence */
2608 0, /* tp_as_mapping */
2609 0, /* tp_hash */
2610 0, /* tp_call */
2611 0, /* tp_str */
2612 PyObject_GenericGetAttr, /* tp_getattro */
2613 0, /* tp_setattro */
2614 0, /* tp_as_buffer */
2615 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2616 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002617 combinations_doc, /* tp_doc */
2618 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002619 0, /* tp_clear */
2620 0, /* tp_richcompare */
2621 0, /* tp_weaklistoffset */
2622 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002623 (iternextfunc)combinations_next, /* tp_iternext */
2624 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002625 0, /* tp_members */
2626 0, /* tp_getset */
2627 0, /* tp_base */
2628 0, /* tp_dict */
2629 0, /* tp_descr_get */
2630 0, /* tp_descr_set */
2631 0, /* tp_dictoffset */
2632 0, /* tp_init */
2633 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002634 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002635 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002636};
2637
2638
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002639/* combinations with replacement object *******************************************/
2640
2641/* Equivalent to:
2642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002643 def combinations_with_replacement(iterable, r):
2644 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2645 # number items returned: (n+r-1)! / r! / (n-1)!
2646 pool = tuple(iterable)
2647 n = len(pool)
2648 indices = [0] * r
2649 yield tuple(pool[i] for i in indices)
2650 while 1:
2651 for i in reversed(range(r)):
2652 if indices[i] != n - 1:
2653 break
2654 else:
2655 return
2656 indices[i:] = [indices[i] + 1] * (r - i)
2657 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002659 def combinations_with_replacement2(iterable, r):
2660 'Alternate version that filters from product()'
2661 pool = tuple(iterable)
2662 n = len(pool)
2663 for indices in product(range(n), repeat=r):
2664 if sorted(indices) == list(indices):
2665 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002666*/
2667typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002668 PyObject_HEAD
2669 PyObject *pool; /* input converted to a tuple */
2670 Py_ssize_t *indices; /* one index per result element */
2671 PyObject *result; /* most recently returned result tuple */
2672 Py_ssize_t r; /* size of result tuple */
2673 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002674} cwrobject;
2675
2676static PyTypeObject cwr_type;
2677
2678static PyObject *
2679cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2680{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002681 cwrobject *co;
2682 Py_ssize_t n;
2683 Py_ssize_t r;
2684 PyObject *pool = NULL;
2685 PyObject *iterable = NULL;
2686 Py_ssize_t *indices = NULL;
2687 Py_ssize_t i;
2688 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002690 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2691 &iterable, &r))
2692 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002693
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002694 pool = PySequence_Tuple(iterable);
2695 if (pool == NULL)
2696 goto error;
2697 n = PyTuple_GET_SIZE(pool);
2698 if (r < 0) {
2699 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2700 goto error;
2701 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002702
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002703 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2704 if (indices == NULL) {
2705 PyErr_NoMemory();
2706 goto error;
2707 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002708
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002709 for (i=0 ; i<r ; i++)
2710 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002711
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002712 /* create cwrobject structure */
2713 co = (cwrobject *)type->tp_alloc(type, 0);
2714 if (co == NULL)
2715 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002716
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002717 co->pool = pool;
2718 co->indices = indices;
2719 co->result = NULL;
2720 co->r = r;
2721 co->stopped = !n && r;
2722
2723 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002724
2725error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002726 if (indices != NULL)
2727 PyMem_Free(indices);
2728 Py_XDECREF(pool);
2729 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002730}
2731
2732static void
2733cwr_dealloc(cwrobject *co)
2734{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002735 PyObject_GC_UnTrack(co);
2736 Py_XDECREF(co->pool);
2737 Py_XDECREF(co->result);
2738 if (co->indices != NULL)
2739 PyMem_Free(co->indices);
2740 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002741}
2742
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002743static PyObject *
2744cwr_sizeof(cwrobject *co, void *unused)
2745{
2746 Py_ssize_t res;
2747
2748 res = sizeof(cwrobject);
2749 res += co->r * sizeof(Py_ssize_t);
2750 return PyLong_FromSsize_t(res);
2751}
2752
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002753static int
2754cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2755{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 Py_VISIT(co->pool);
2757 Py_VISIT(co->result);
2758 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002759}
2760
2761static PyObject *
2762cwr_next(cwrobject *co)
2763{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002764 PyObject *elem;
2765 PyObject *oldelem;
2766 PyObject *pool = co->pool;
2767 Py_ssize_t *indices = co->indices;
2768 PyObject *result = co->result;
2769 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2770 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002771 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002772
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002773 if (co->stopped)
2774 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002775
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002776 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002777 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002778 result = PyTuple_New(r);
2779 if (result == NULL)
2780 goto empty;
2781 co->result = result;
Tim Peters9edb1682013-09-03 11:49:31 -05002782 elem = PyTuple_GET_ITEM(pool, 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002783 for (i=0; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002784 assert(indices[i] == 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002785 Py_INCREF(elem);
2786 PyTuple_SET_ITEM(result, i, elem);
2787 }
2788 } else {
2789 /* Copy the previous result tuple or re-use it if available */
2790 if (Py_REFCNT(result) > 1) {
2791 PyObject *old_result = result;
2792 result = PyTuple_New(r);
2793 if (result == NULL)
2794 goto empty;
2795 co->result = result;
2796 for (i=0; i<r ; i++) {
2797 elem = PyTuple_GET_ITEM(old_result, i);
2798 Py_INCREF(elem);
2799 PyTuple_SET_ITEM(result, i, elem);
2800 }
2801 Py_DECREF(old_result);
2802 }
2803 /* Now, we've got the only copy so we can update it in-place CPython's
2804 empty tuple is a singleton and cached in PyTuple's freelist. */
2805 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002806
Tim Peters9edb1682013-09-03 11:49:31 -05002807 /* Scan indices right-to-left until finding one that is not
2808 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2810 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002813 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002814 if (i < 0)
2815 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002816
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002817 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002818 maximum. Then set all to the right to the same value. */
2819 index = indices[i] + 1;
2820 assert(index < n);
2821 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002823 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002824 Py_INCREF(elem);
2825 oldelem = PyTuple_GET_ITEM(result, i);
2826 PyTuple_SET_ITEM(result, i, elem);
2827 Py_DECREF(oldelem);
2828 }
2829 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002831 Py_INCREF(result);
2832 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002833
2834empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002835 co->stopped = 1;
2836 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002837}
2838
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002839static PyObject *
2840cwr_reduce(cwrobject *lz)
2841{
2842 if (lz->result == NULL) {
2843 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2844 } else if (lz->stopped) {
2845 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2846 } else {
2847 PyObject *indices;
2848 Py_ssize_t i;
2849
2850 /* we must pickle the indices and use them for setstate */
2851 indices = PyTuple_New(lz->r);
2852 if (!indices)
2853 return NULL;
2854 for (i=0; i<lz->r; i++)
2855 {
2856 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2857 if (!index) {
2858 Py_DECREF(indices);
2859 return NULL;
2860 }
2861 PyTuple_SET_ITEM(indices, i, index);
2862 }
2863
2864 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2865 }
2866}
2867
2868static PyObject *
2869cwr_setstate(cwrobject *lz, PyObject *state)
2870{
2871 PyObject *result;
2872 Py_ssize_t n, i;
2873
2874 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2875 {
2876 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2877 return NULL;
2878 }
2879
2880 n = PyTuple_GET_SIZE(lz->pool);
2881 for (i=0; i<lz->r; i++)
2882 {
2883 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2884 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2885 if (index < 0 && PyErr_Occurred())
2886 return NULL; /* not an integer */
2887 /* clamp the index */
2888 if (index < 0)
2889 index = 0;
2890 else if (index > n-1)
2891 index = n-1;
2892 lz->indices[i] = index;
2893 }
2894 result = PyTuple_New(lz->r);
2895 if (result == NULL)
2896 return NULL;
2897 for (i=0; i<lz->r; i++) {
2898 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2899 Py_INCREF(element);
2900 PyTuple_SET_ITEM(result, i, element);
2901 }
2902 Py_CLEAR(lz->result);
2903 lz->result = result;
2904 Py_RETURN_NONE;
2905}
2906
2907static PyMethodDef cwr_methods[] = {
2908 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2909 reduce_doc},
2910 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2911 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002912 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2913 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002914 {NULL, NULL} /* sentinel */
2915};
2916
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002917PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002918"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002919\n\
2920Return successive r-length combinations of elements in the iterable\n\
2921allowing individual elements to have successive repeats.\n\
2922combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2923
2924static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002925 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002926 "itertools.combinations_with_replacement", /* tp_name */
2927 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002928 0, /* tp_itemsize */
2929 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002930 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002931 0, /* tp_print */
2932 0, /* tp_getattr */
2933 0, /* tp_setattr */
2934 0, /* tp_reserved */
2935 0, /* tp_repr */
2936 0, /* tp_as_number */
2937 0, /* tp_as_sequence */
2938 0, /* tp_as_mapping */
2939 0, /* tp_hash */
2940 0, /* tp_call */
2941 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002942 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002943 0, /* tp_setattro */
2944 0, /* tp_as_buffer */
2945 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002946 Py_TPFLAGS_BASETYPE, /* tp_flags */
2947 cwr_doc, /* tp_doc */
2948 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002949 0, /* tp_clear */
2950 0, /* tp_richcompare */
2951 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002952 PyObject_SelfIter, /* tp_iter */
2953 (iternextfunc)cwr_next, /* tp_iternext */
2954 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002955 0, /* tp_members */
2956 0, /* tp_getset */
2957 0, /* tp_base */
2958 0, /* tp_dict */
2959 0, /* tp_descr_get */
2960 0, /* tp_descr_set */
2961 0, /* tp_dictoffset */
2962 0, /* tp_init */
2963 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002964 cwr_new, /* tp_new */
2965 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002966};
2967
2968
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002969/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002970
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002971def permutations(iterable, r=None):
2972 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2973 pool = tuple(iterable)
2974 n = len(pool)
2975 r = n if r is None else r
2976 indices = range(n)
2977 cycles = range(n-r+1, n+1)[::-1]
2978 yield tuple(pool[i] for i in indices[:r])
2979 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002980 for i in reversed(range(r)):
2981 cycles[i] -= 1
2982 if cycles[i] == 0:
2983 indices[i:] = indices[i+1:] + indices[i:i+1]
2984 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002985 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002986 j = cycles[i]
2987 indices[i], indices[-j] = indices[-j], indices[i]
2988 yield tuple(pool[i] for i in indices[:r])
2989 break
2990 else:
2991 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002992*/
2993
2994typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002995 PyObject_HEAD
2996 PyObject *pool; /* input converted to a tuple */
2997 Py_ssize_t *indices; /* one index per element in the pool */
2998 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2999 PyObject *result; /* most recently returned result tuple */
3000 Py_ssize_t r; /* size of result tuple */
3001 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003002} permutationsobject;
3003
3004static PyTypeObject permutations_type;
3005
3006static PyObject *
3007permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3008{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003009 permutationsobject *po;
3010 Py_ssize_t n;
3011 Py_ssize_t r;
3012 PyObject *robj = Py_None;
3013 PyObject *pool = NULL;
3014 PyObject *iterable = NULL;
3015 Py_ssize_t *indices = NULL;
3016 Py_ssize_t *cycles = NULL;
3017 Py_ssize_t i;
3018 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003019
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003020 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3021 &iterable, &robj))
3022 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003023
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003024 pool = PySequence_Tuple(iterable);
3025 if (pool == NULL)
3026 goto error;
3027 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003028
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003029 r = n;
3030 if (robj != Py_None) {
3031 if (!PyLong_Check(robj)) {
3032 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3033 goto error;
3034 }
3035 r = PyLong_AsSsize_t(robj);
3036 if (r == -1 && PyErr_Occurred())
3037 goto error;
3038 }
3039 if (r < 0) {
3040 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3041 goto error;
3042 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003043
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003044 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
3045 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
3046 if (indices == NULL || cycles == NULL) {
3047 PyErr_NoMemory();
3048 goto error;
3049 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003050
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003051 for (i=0 ; i<n ; i++)
3052 indices[i] = i;
3053 for (i=0 ; i<r ; i++)
3054 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003055
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003056 /* create permutationsobject structure */
3057 po = (permutationsobject *)type->tp_alloc(type, 0);
3058 if (po == NULL)
3059 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003060
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 po->pool = pool;
3062 po->indices = indices;
3063 po->cycles = cycles;
3064 po->result = NULL;
3065 po->r = r;
3066 po->stopped = r > n ? 1 : 0;
3067
3068 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003069
3070error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003071 if (indices != NULL)
3072 PyMem_Free(indices);
3073 if (cycles != NULL)
3074 PyMem_Free(cycles);
3075 Py_XDECREF(pool);
3076 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003077}
3078
3079static void
3080permutations_dealloc(permutationsobject *po)
3081{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003082 PyObject_GC_UnTrack(po);
3083 Py_XDECREF(po->pool);
3084 Py_XDECREF(po->result);
3085 PyMem_Free(po->indices);
3086 PyMem_Free(po->cycles);
3087 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003088}
3089
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003090static PyObject *
3091permutations_sizeof(permutationsobject *po, void *unused)
3092{
3093 Py_ssize_t res;
3094
3095 res = sizeof(permutationsobject);
3096 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3097 res += po->r * sizeof(Py_ssize_t);
3098 return PyLong_FromSsize_t(res);
3099}
3100
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003101static int
3102permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3103{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003104 Py_VISIT(po->pool);
3105 Py_VISIT(po->result);
3106 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003107}
3108
3109static PyObject *
3110permutations_next(permutationsobject *po)
3111{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003112 PyObject *elem;
3113 PyObject *oldelem;
3114 PyObject *pool = po->pool;
3115 Py_ssize_t *indices = po->indices;
3116 Py_ssize_t *cycles = po->cycles;
3117 PyObject *result = po->result;
3118 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3119 Py_ssize_t r = po->r;
3120 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 if (po->stopped)
3123 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003124
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003125 if (result == NULL) {
3126 /* On the first pass, initialize result tuple using the indices */
3127 result = PyTuple_New(r);
3128 if (result == NULL)
3129 goto empty;
3130 po->result = result;
3131 for (i=0; i<r ; i++) {
3132 index = indices[i];
3133 elem = PyTuple_GET_ITEM(pool, index);
3134 Py_INCREF(elem);
3135 PyTuple_SET_ITEM(result, i, elem);
3136 }
3137 } else {
3138 if (n == 0)
3139 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003141 /* Copy the previous result tuple or re-use it if available */
3142 if (Py_REFCNT(result) > 1) {
3143 PyObject *old_result = result;
3144 result = PyTuple_New(r);
3145 if (result == NULL)
3146 goto empty;
3147 po->result = result;
3148 for (i=0; i<r ; i++) {
3149 elem = PyTuple_GET_ITEM(old_result, i);
3150 Py_INCREF(elem);
3151 PyTuple_SET_ITEM(result, i, elem);
3152 }
3153 Py_DECREF(old_result);
3154 }
3155 /* Now, we've got the only copy so we can update it in-place */
3156 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003157
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003158 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3159 for (i=r-1 ; i>=0 ; i--) {
3160 cycles[i] -= 1;
3161 if (cycles[i] == 0) {
3162 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3163 index = indices[i];
3164 for (j=i ; j<n-1 ; j++)
3165 indices[j] = indices[j+1];
3166 indices[n-1] = index;
3167 cycles[i] = n - i;
3168 } else {
3169 j = cycles[i];
3170 index = indices[i];
3171 indices[i] = indices[n-j];
3172 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003173
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003174 for (k=i; k<r ; k++) {
3175 /* start with i, the leftmost element that changed */
3176 /* yield tuple(pool[k] for k in indices[:r]) */
3177 index = indices[k];
3178 elem = PyTuple_GET_ITEM(pool, index);
3179 Py_INCREF(elem);
3180 oldelem = PyTuple_GET_ITEM(result, k);
3181 PyTuple_SET_ITEM(result, k, elem);
3182 Py_DECREF(oldelem);
3183 }
3184 break;
3185 }
3186 }
3187 /* If i is negative, then the cycles have all
3188 rolled-over and we're done. */
3189 if (i < 0)
3190 goto empty;
3191 }
3192 Py_INCREF(result);
3193 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003194
3195empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003196 po->stopped = 1;
3197 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003198}
3199
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003200static PyObject *
3201permutations_reduce(permutationsobject *po)
3202{
3203 if (po->result == NULL) {
3204 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3205 } else if (po->stopped) {
3206 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3207 } else {
3208 PyObject *indices=NULL, *cycles=NULL;
3209 Py_ssize_t n, i;
3210
3211 /* we must pickle the indices and cycles and use them for setstate */
3212 n = PyTuple_GET_SIZE(po->pool);
3213 indices = PyTuple_New(n);
3214 if (indices == NULL)
3215 goto err;
3216 for (i=0; i<n; i++){
3217 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3218 if (!index)
3219 goto err;
3220 PyTuple_SET_ITEM(indices, i, index);
3221 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003222
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003223 cycles = PyTuple_New(po->r);
3224 if (cycles == NULL)
3225 goto err;
3226 for (i=0; i<po->r; i++)
3227 {
3228 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3229 if (!index)
3230 goto err;
3231 PyTuple_SET_ITEM(cycles, i, index);
3232 }
3233 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3234 po->pool, po->r,
3235 indices, cycles);
3236 err:
3237 Py_XDECREF(indices);
3238 Py_XDECREF(cycles);
3239 return NULL;
3240 }
3241}
3242
3243static PyObject *
3244permutations_setstate(permutationsobject *po, PyObject *state)
3245{
3246 PyObject *indices, *cycles, *result;
3247 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003248
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003249 if (!PyArg_ParseTuple(state, "O!O!",
3250 &PyTuple_Type, &indices,
3251 &PyTuple_Type, &cycles))
3252 return NULL;
3253
3254 n = PyTuple_GET_SIZE(po->pool);
3255 if (PyTuple_GET_SIZE(indices) != n ||
3256 PyTuple_GET_SIZE(cycles) != po->r)
3257 {
3258 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3259 return NULL;
3260 }
3261
3262 for (i=0; i<n; i++)
3263 {
3264 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3265 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3266 if (index < 0 && PyErr_Occurred())
3267 return NULL; /* not an integer */
3268 /* clamp the index */
3269 if (index < 0)
3270 index = 0;
3271 else if (index > n-1)
3272 index = n-1;
3273 po->indices[i] = index;
3274 }
3275
3276 for (i=0; i<po->r; i++)
3277 {
3278 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3279 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3280 if (index < 0 && PyErr_Occurred())
3281 return NULL; /* not an integer */
3282 if (index < 1)
3283 index = 1;
3284 else if (index > n-i)
3285 index = n-i;
3286 po->cycles[i] = index;
3287 }
3288 result = PyTuple_New(po->r);
3289 if (result == NULL)
3290 return NULL;
3291 for (i=0; i<po->r; i++) {
3292 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3293 Py_INCREF(element);
3294 PyTuple_SET_ITEM(result, i, element);
3295 }
3296 Py_CLEAR(po->result);
3297 po->result = result;
3298 Py_RETURN_NONE;
3299}
3300
3301static PyMethodDef permuations_methods[] = {
3302 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3303 reduce_doc},
3304 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3305 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003306 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3307 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003308 {NULL, NULL} /* sentinel */
3309};
3310
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003311PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003312"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003313\n\
3314Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003315permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003316
3317static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003318 PyVarObject_HEAD_INIT(NULL, 0)
3319 "itertools.permutations", /* tp_name */
3320 sizeof(permutationsobject), /* tp_basicsize */
3321 0, /* tp_itemsize */
3322 /* methods */
3323 (destructor)permutations_dealloc, /* tp_dealloc */
3324 0, /* tp_print */
3325 0, /* tp_getattr */
3326 0, /* tp_setattr */
3327 0, /* tp_reserved */
3328 0, /* tp_repr */
3329 0, /* tp_as_number */
3330 0, /* tp_as_sequence */
3331 0, /* tp_as_mapping */
3332 0, /* tp_hash */
3333 0, /* tp_call */
3334 0, /* tp_str */
3335 PyObject_GenericGetAttr, /* tp_getattro */
3336 0, /* tp_setattro */
3337 0, /* tp_as_buffer */
3338 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3339 Py_TPFLAGS_BASETYPE, /* tp_flags */
3340 permutations_doc, /* tp_doc */
3341 (traverseproc)permutations_traverse, /* tp_traverse */
3342 0, /* tp_clear */
3343 0, /* tp_richcompare */
3344 0, /* tp_weaklistoffset */
3345 PyObject_SelfIter, /* tp_iter */
3346 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003347 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003348 0, /* tp_members */
3349 0, /* tp_getset */
3350 0, /* tp_base */
3351 0, /* tp_dict */
3352 0, /* tp_descr_get */
3353 0, /* tp_descr_set */
3354 0, /* tp_dictoffset */
3355 0, /* tp_init */
3356 0, /* tp_alloc */
3357 permutations_new, /* tp_new */
3358 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003359};
3360
Raymond Hettinger482ba772010-12-01 22:48:00 +00003361/* accumulate object ************************************************************/
3362
3363typedef struct {
3364 PyObject_HEAD
3365 PyObject *total;
3366 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003367 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003368} accumulateobject;
3369
3370static PyTypeObject accumulate_type;
3371
3372static PyObject *
3373accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3374{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003375 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003376 PyObject *iterable;
3377 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003378 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003379 accumulateobject *lz;
3380
Raymond Hettinger5d446132011-03-27 18:52:10 -07003381 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3382 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003383 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003384
3385 /* Get iterator. */
3386 it = PyObject_GetIter(iterable);
3387 if (it == NULL)
3388 return NULL;
3389
Raymond Hettinger482ba772010-12-01 22:48:00 +00003390 /* create accumulateobject structure */
3391 lz = (accumulateobject *)type->tp_alloc(type, 0);
3392 if (lz == NULL) {
3393 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003394 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003395 }
3396
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003397 if (binop != Py_None) {
3398 Py_XINCREF(binop);
3399 lz->binop = binop;
3400 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003401 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003402 lz->it = it;
3403 return (PyObject *)lz;
3404}
3405
3406static void
3407accumulate_dealloc(accumulateobject *lz)
3408{
3409 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003410 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003411 Py_XDECREF(lz->total);
3412 Py_XDECREF(lz->it);
3413 Py_TYPE(lz)->tp_free(lz);
3414}
3415
3416static int
3417accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3418{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003419 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003420 Py_VISIT(lz->it);
3421 Py_VISIT(lz->total);
3422 return 0;
3423}
3424
3425static PyObject *
3426accumulate_next(accumulateobject *lz)
3427{
3428 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003429
Raymond Hettinger482ba772010-12-01 22:48:00 +00003430 val = PyIter_Next(lz->it);
3431 if (val == NULL)
3432 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003433
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003434 if (lz->total == NULL) {
3435 Py_INCREF(val);
3436 lz->total = val;
3437 return lz->total;
3438 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003439
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003440 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003441 newtotal = PyNumber_Add(lz->total, val);
3442 else
3443 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003444 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003445 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003446 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003447
3448 oldtotal = lz->total;
3449 lz->total = newtotal;
3450 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003451
Raymond Hettinger482ba772010-12-01 22:48:00 +00003452 Py_INCREF(newtotal);
3453 return newtotal;
3454}
3455
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003456static PyObject *
3457accumulate_reduce(accumulateobject *lz)
3458{
3459 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3460 lz->it, lz->binop?lz->binop:Py_None,
3461 lz->total?lz->total:Py_None);
3462 }
3463
3464static PyObject *
3465accumulate_setstate(accumulateobject *lz, PyObject *state)
3466{
3467 Py_CLEAR(lz->total);
3468 lz->total = state;
3469 Py_INCREF(lz->total);
3470 Py_RETURN_NONE;
3471}
3472
3473static PyMethodDef accumulate_methods[] = {
3474 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3475 reduce_doc},
3476 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3477 setstate_doc},
3478 {NULL, NULL} /* sentinel */
3479};
3480
Raymond Hettinger482ba772010-12-01 22:48:00 +00003481PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003482"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003483\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003484Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003485
3486static PyTypeObject accumulate_type = {
3487 PyVarObject_HEAD_INIT(NULL, 0)
3488 "itertools.accumulate", /* tp_name */
3489 sizeof(accumulateobject), /* tp_basicsize */
3490 0, /* tp_itemsize */
3491 /* methods */
3492 (destructor)accumulate_dealloc, /* tp_dealloc */
3493 0, /* tp_print */
3494 0, /* tp_getattr */
3495 0, /* tp_setattr */
3496 0, /* tp_reserved */
3497 0, /* tp_repr */
3498 0, /* tp_as_number */
3499 0, /* tp_as_sequence */
3500 0, /* tp_as_mapping */
3501 0, /* tp_hash */
3502 0, /* tp_call */
3503 0, /* tp_str */
3504 PyObject_GenericGetAttr, /* tp_getattro */
3505 0, /* tp_setattro */
3506 0, /* tp_as_buffer */
3507 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3508 Py_TPFLAGS_BASETYPE, /* tp_flags */
3509 accumulate_doc, /* tp_doc */
3510 (traverseproc)accumulate_traverse, /* tp_traverse */
3511 0, /* tp_clear */
3512 0, /* tp_richcompare */
3513 0, /* tp_weaklistoffset */
3514 PyObject_SelfIter, /* tp_iter */
3515 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003516 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003517 0, /* tp_members */
3518 0, /* tp_getset */
3519 0, /* tp_base */
3520 0, /* tp_dict */
3521 0, /* tp_descr_get */
3522 0, /* tp_descr_set */
3523 0, /* tp_dictoffset */
3524 0, /* tp_init */
3525 0, /* tp_alloc */
3526 accumulate_new, /* tp_new */
3527 PyObject_GC_Del, /* tp_free */
3528};
3529
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003530
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003531/* compress object ************************************************************/
3532
3533/* Equivalent to:
3534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003535 def compress(data, selectors):
3536 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3537 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003538*/
3539
3540typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003541 PyObject_HEAD
3542 PyObject *data;
3543 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003544} compressobject;
3545
3546static PyTypeObject compress_type;
3547
3548static PyObject *
3549compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3550{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003551 PyObject *seq1, *seq2;
3552 PyObject *data=NULL, *selectors=NULL;
3553 compressobject *lz;
3554 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003555
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003556 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3557 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003558
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003559 data = PyObject_GetIter(seq1);
3560 if (data == NULL)
3561 goto fail;
3562 selectors = PyObject_GetIter(seq2);
3563 if (selectors == NULL)
3564 goto fail;
3565
3566 /* create compressobject structure */
3567 lz = (compressobject *)type->tp_alloc(type, 0);
3568 if (lz == NULL)
3569 goto fail;
3570 lz->data = data;
3571 lz->selectors = selectors;
3572 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003573
3574fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003575 Py_XDECREF(data);
3576 Py_XDECREF(selectors);
3577 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003578}
3579
3580static void
3581compress_dealloc(compressobject *lz)
3582{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003583 PyObject_GC_UnTrack(lz);
3584 Py_XDECREF(lz->data);
3585 Py_XDECREF(lz->selectors);
3586 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003587}
3588
3589static int
3590compress_traverse(compressobject *lz, visitproc visit, void *arg)
3591{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003592 Py_VISIT(lz->data);
3593 Py_VISIT(lz->selectors);
3594 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003595}
3596
3597static PyObject *
3598compress_next(compressobject *lz)
3599{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003600 PyObject *data = lz->data, *selectors = lz->selectors;
3601 PyObject *datum, *selector;
3602 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3603 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3604 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003605
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003606 while (1) {
3607 /* Steps: get datum, get selector, evaluate selector.
3608 Order is important (to match the pure python version
3609 in terms of which input gets a chance to raise an
3610 exception first).
3611 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003612
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003613 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003614 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003615 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003616
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003617 selector = selectornext(selectors);
3618 if (selector == NULL) {
3619 Py_DECREF(datum);
3620 return NULL;
3621 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003622
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003623 ok = PyObject_IsTrue(selector);
3624 Py_DECREF(selector);
3625 if (ok == 1)
3626 return datum;
3627 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003628 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003629 return NULL;
3630 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003631}
3632
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003633static PyObject *
3634compress_reduce(compressobject *lz)
3635{
3636 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3637 lz->data, lz->selectors);
3638 }
3639
3640static PyMethodDef compress_methods[] = {
3641 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3642 reduce_doc},
3643 {NULL, NULL} /* sentinel */
3644};
3645
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003646PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003647"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003648\n\
3649Return data elements corresponding to true selector elements.\n\
3650Forms a shorter iterator from selected data elements using the\n\
3651selectors to choose the data elements.");
3652
3653static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003654 PyVarObject_HEAD_INIT(NULL, 0)
3655 "itertools.compress", /* tp_name */
3656 sizeof(compressobject), /* tp_basicsize */
3657 0, /* tp_itemsize */
3658 /* methods */
3659 (destructor)compress_dealloc, /* tp_dealloc */
3660 0, /* tp_print */
3661 0, /* tp_getattr */
3662 0, /* tp_setattr */
3663 0, /* tp_reserved */
3664 0, /* tp_repr */
3665 0, /* tp_as_number */
3666 0, /* tp_as_sequence */
3667 0, /* tp_as_mapping */
3668 0, /* tp_hash */
3669 0, /* tp_call */
3670 0, /* tp_str */
3671 PyObject_GenericGetAttr, /* tp_getattro */
3672 0, /* tp_setattro */
3673 0, /* tp_as_buffer */
3674 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3675 Py_TPFLAGS_BASETYPE, /* tp_flags */
3676 compress_doc, /* tp_doc */
3677 (traverseproc)compress_traverse, /* tp_traverse */
3678 0, /* tp_clear */
3679 0, /* tp_richcompare */
3680 0, /* tp_weaklistoffset */
3681 PyObject_SelfIter, /* tp_iter */
3682 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003683 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003684 0, /* tp_members */
3685 0, /* tp_getset */
3686 0, /* tp_base */
3687 0, /* tp_dict */
3688 0, /* tp_descr_get */
3689 0, /* tp_descr_set */
3690 0, /* tp_dictoffset */
3691 0, /* tp_init */
3692 0, /* tp_alloc */
3693 compress_new, /* tp_new */
3694 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003695};
3696
3697
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003698/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003699
3700typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003701 PyObject_HEAD
3702 PyObject *func;
3703 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003704} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003705
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003706static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003707
3708static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003709filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003710{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003711 PyObject *func, *seq;
3712 PyObject *it;
3713 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003714
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003715 if (type == &filterfalse_type &&
3716 !_PyArg_NoKeywords("filterfalse()", kwds))
3717 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003718
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003719 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3720 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003721
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003722 /* Get iterator. */
3723 it = PyObject_GetIter(seq);
3724 if (it == NULL)
3725 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003726
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003727 /* create filterfalseobject structure */
3728 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3729 if (lz == NULL) {
3730 Py_DECREF(it);
3731 return NULL;
3732 }
3733 Py_INCREF(func);
3734 lz->func = func;
3735 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003736
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003737 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003738}
3739
3740static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003741filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003742{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003743 PyObject_GC_UnTrack(lz);
3744 Py_XDECREF(lz->func);
3745 Py_XDECREF(lz->it);
3746 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003747}
3748
3749static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003750filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003751{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003752 Py_VISIT(lz->it);
3753 Py_VISIT(lz->func);
3754 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003755}
3756
3757static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003758filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003759{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003760 PyObject *item;
3761 PyObject *it = lz->it;
3762 long ok;
3763 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003764
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003765 iternext = *Py_TYPE(it)->tp_iternext;
3766 for (;;) {
3767 item = iternext(it);
3768 if (item == NULL)
3769 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003771 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3772 ok = PyObject_IsTrue(item);
3773 } else {
3774 PyObject *good;
3775 good = PyObject_CallFunctionObjArgs(lz->func,
3776 item, NULL);
3777 if (good == NULL) {
3778 Py_DECREF(item);
3779 return NULL;
3780 }
3781 ok = PyObject_IsTrue(good);
3782 Py_DECREF(good);
3783 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003784 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003785 return item;
3786 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003787 if (ok < 0)
3788 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003789 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003790}
3791
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003792static PyObject *
3793filterfalse_reduce(filterfalseobject *lz)
3794{
3795 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3796 lz->func, lz->it);
3797 }
3798
3799static PyMethodDef filterfalse_methods[] = {
3800 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3801 reduce_doc},
3802 {NULL, NULL} /* sentinel */
3803};
3804
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003805PyDoc_STRVAR(filterfalse_doc,
3806"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003807\n\
3808Return those items of sequence for which function(item) is false.\n\
3809If function is None, return the items that are false.");
3810
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003811static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003812 PyVarObject_HEAD_INIT(NULL, 0)
3813 "itertools.filterfalse", /* tp_name */
3814 sizeof(filterfalseobject), /* tp_basicsize */
3815 0, /* tp_itemsize */
3816 /* methods */
3817 (destructor)filterfalse_dealloc, /* tp_dealloc */
3818 0, /* tp_print */
3819 0, /* tp_getattr */
3820 0, /* tp_setattr */
3821 0, /* tp_reserved */
3822 0, /* tp_repr */
3823 0, /* tp_as_number */
3824 0, /* tp_as_sequence */
3825 0, /* tp_as_mapping */
3826 0, /* tp_hash */
3827 0, /* tp_call */
3828 0, /* tp_str */
3829 PyObject_GenericGetAttr, /* tp_getattro */
3830 0, /* tp_setattro */
3831 0, /* tp_as_buffer */
3832 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3833 Py_TPFLAGS_BASETYPE, /* tp_flags */
3834 filterfalse_doc, /* tp_doc */
3835 (traverseproc)filterfalse_traverse, /* tp_traverse */
3836 0, /* tp_clear */
3837 0, /* tp_richcompare */
3838 0, /* tp_weaklistoffset */
3839 PyObject_SelfIter, /* tp_iter */
3840 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003841 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003842 0, /* tp_members */
3843 0, /* tp_getset */
3844 0, /* tp_base */
3845 0, /* tp_dict */
3846 0, /* tp_descr_get */
3847 0, /* tp_descr_set */
3848 0, /* tp_dictoffset */
3849 0, /* tp_init */
3850 0, /* tp_alloc */
3851 filterfalse_new, /* tp_new */
3852 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003853};
3854
3855
3856/* count object ************************************************************/
3857
3858typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003859 PyObject_HEAD
3860 Py_ssize_t cnt;
3861 PyObject *long_cnt;
3862 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003863} countobject;
3864
Raymond Hettinger30729212009-02-12 06:28:27 +00003865/* Counting logic and invariants:
3866
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003867fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003868
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003869 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3870 Advances with: cnt += 1
3871 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003872
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003873slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003875 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3876 All counting is done with python objects (no overflows or underflows).
3877 Advances with: long_cnt += long_step
3878 Step may be zero -- effectively a slow version of repeat(cnt).
3879 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003880*/
3881
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003882static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003883
3884static PyObject *
3885count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3886{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003887 countobject *lz;
3888 int slow_mode = 0;
3889 Py_ssize_t cnt = 0;
3890 PyObject *long_cnt = NULL;
3891 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003892 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003893 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003895 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3896 kwlist, &long_cnt, &long_step))
3897 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003898
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003899 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3900 (long_step != NULL && !PyNumber_Check(long_step))) {
3901 PyErr_SetString(PyExc_TypeError, "a number is required");
3902 return NULL;
3903 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003904
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003905 if (long_cnt != NULL) {
3906 cnt = PyLong_AsSsize_t(long_cnt);
3907 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3908 PyErr_Clear();
3909 slow_mode = 1;
3910 }
3911 Py_INCREF(long_cnt);
3912 } else {
3913 cnt = 0;
3914 long_cnt = PyLong_FromLong(0);
3915 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003916
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003917 /* If not specified, step defaults to 1 */
3918 if (long_step == NULL) {
3919 long_step = PyLong_FromLong(1);
3920 if (long_step == NULL) {
3921 Py_DECREF(long_cnt);
3922 return NULL;
3923 }
3924 } else
3925 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003927 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003928
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003929 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003930 step = PyLong_AsLong(long_step);
3931 if (step != 1) {
3932 slow_mode = 1;
3933 if (step == -1 && PyErr_Occurred())
3934 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003936
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 if (slow_mode)
3938 cnt = PY_SSIZE_T_MAX;
3939 else
3940 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003941
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003942 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3943 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3944 assert(slow_mode ||
3945 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 /* create countobject structure */
3948 lz = (countobject *)type->tp_alloc(type, 0);
3949 if (lz == NULL) {
3950 Py_XDECREF(long_cnt);
3951 return NULL;
3952 }
3953 lz->cnt = cnt;
3954 lz->long_cnt = long_cnt;
3955 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003956
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003958}
3959
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003960static void
3961count_dealloc(countobject *lz)
3962{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003963 PyObject_GC_UnTrack(lz);
3964 Py_XDECREF(lz->long_cnt);
3965 Py_XDECREF(lz->long_step);
3966 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003967}
3968
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003969static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003970count_traverse(countobject *lz, visitproc visit, void *arg)
3971{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003972 Py_VISIT(lz->long_cnt);
3973 Py_VISIT(lz->long_step);
3974 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003975}
3976
3977static PyObject *
3978count_nextlong(countobject *lz)
3979{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003980 PyObject *long_cnt;
3981 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003982
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003983 long_cnt = lz->long_cnt;
3984 if (long_cnt == NULL) {
3985 /* Switch to slow_mode */
3986 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3987 if (long_cnt == NULL)
3988 return NULL;
3989 }
3990 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003991
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003992 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3993 if (stepped_up == NULL)
3994 return NULL;
3995 lz->long_cnt = stepped_up;
3996 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003997}
3998
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003999static PyObject *
4000count_next(countobject *lz)
4001{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004002 if (lz->cnt == PY_SSIZE_T_MAX)
4003 return count_nextlong(lz);
4004 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004005}
4006
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004007static PyObject *
4008count_repr(countobject *lz)
4009{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004010 if (lz->cnt != PY_SSIZE_T_MAX)
4011 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004012
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004013 if (PyLong_Check(lz->long_step)) {
4014 long step = PyLong_AsLong(lz->long_step);
4015 if (step == -1 && PyErr_Occurred()) {
4016 PyErr_Clear();
4017 }
4018 if (step == 1) {
4019 /* Don't display step when it is an integer equal to 1 */
4020 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4021 }
4022 }
4023 return PyUnicode_FromFormat("count(%R, %R)",
4024 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004025}
4026
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004027static PyObject *
4028count_reduce(countobject *lz)
4029{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004030 if (lz->cnt == PY_SSIZE_T_MAX)
4031 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4032 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004033}
4034
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004035static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004036 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004037 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004038 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004039};
4040
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004041PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004042 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004043\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004044Return a count object whose .__next__() method returns consecutive values.\n\
4045Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004046 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004047 x = firstval\n\
4048 while 1:\n\
4049 yield x\n\
4050 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004051
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004052static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004053 PyVarObject_HEAD_INIT(NULL, 0)
4054 "itertools.count", /* tp_name */
4055 sizeof(countobject), /* tp_basicsize */
4056 0, /* tp_itemsize */
4057 /* methods */
4058 (destructor)count_dealloc, /* tp_dealloc */
4059 0, /* tp_print */
4060 0, /* tp_getattr */
4061 0, /* tp_setattr */
4062 0, /* tp_reserved */
4063 (reprfunc)count_repr, /* tp_repr */
4064 0, /* tp_as_number */
4065 0, /* tp_as_sequence */
4066 0, /* tp_as_mapping */
4067 0, /* tp_hash */
4068 0, /* tp_call */
4069 0, /* tp_str */
4070 PyObject_GenericGetAttr, /* tp_getattro */
4071 0, /* tp_setattro */
4072 0, /* tp_as_buffer */
4073 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4074 Py_TPFLAGS_BASETYPE, /* tp_flags */
4075 count_doc, /* tp_doc */
4076 (traverseproc)count_traverse, /* tp_traverse */
4077 0, /* tp_clear */
4078 0, /* tp_richcompare */
4079 0, /* tp_weaklistoffset */
4080 PyObject_SelfIter, /* tp_iter */
4081 (iternextfunc)count_next, /* tp_iternext */
4082 count_methods, /* tp_methods */
4083 0, /* tp_members */
4084 0, /* tp_getset */
4085 0, /* tp_base */
4086 0, /* tp_dict */
4087 0, /* tp_descr_get */
4088 0, /* tp_descr_set */
4089 0, /* tp_dictoffset */
4090 0, /* tp_init */
4091 0, /* tp_alloc */
4092 count_new, /* tp_new */
4093 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004094};
4095
4096
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004097/* repeat object ************************************************************/
4098
4099typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004100 PyObject_HEAD
4101 PyObject *element;
4102 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004103} repeatobject;
4104
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004105static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004106
4107static PyObject *
4108repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4109{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004110 repeatobject *ro;
4111 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004112 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004113 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004114
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004115 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4116 &element, &cnt))
4117 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004118
Raymond Hettinger97d35552014-06-24 21:36:58 -07004119 if (kwds != NULL)
4120 n_kwds = PyDict_Size(kwds);
4121 /* Does user supply times argument? */
4122 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 cnt = 0;
4124
4125 ro = (repeatobject *)type->tp_alloc(type, 0);
4126 if (ro == NULL)
4127 return NULL;
4128 Py_INCREF(element);
4129 ro->element = element;
4130 ro->cnt = cnt;
4131 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004132}
4133
4134static void
4135repeat_dealloc(repeatobject *ro)
4136{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004137 PyObject_GC_UnTrack(ro);
4138 Py_XDECREF(ro->element);
4139 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004140}
4141
4142static int
4143repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4144{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004145 Py_VISIT(ro->element);
4146 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004147}
4148
4149static PyObject *
4150repeat_next(repeatobject *ro)
4151{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004152 if (ro->cnt == 0)
4153 return NULL;
4154 if (ro->cnt > 0)
4155 ro->cnt--;
4156 Py_INCREF(ro->element);
4157 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004158}
4159
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004160static PyObject *
4161repeat_repr(repeatobject *ro)
4162{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004163 if (ro->cnt == -1)
4164 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4165 else
4166 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4167}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004168
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004169static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004170repeat_len(repeatobject *ro)
4171{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004172 if (ro->cnt == -1) {
4173 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4174 return NULL;
4175 }
4176 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004177}
4178
Armin Rigof5b3e362006-02-11 21:32:43 +00004179PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004180
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004181static PyObject *
4182repeat_reduce(repeatobject *ro)
4183{
4184 /* unpickle this so that a new repeat iterator is constructed with an
4185 * object, then call __setstate__ on it to set cnt
4186 */
4187 if (ro->cnt >= 0)
4188 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4189 else
4190 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4191}
4192
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004193static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004194 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004195 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004197};
4198
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004199PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004200"repeat(object [,times]) -> create an iterator which returns the object\n\
4201for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004202endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004203
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004204static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004205 PyVarObject_HEAD_INIT(NULL, 0)
4206 "itertools.repeat", /* tp_name */
4207 sizeof(repeatobject), /* tp_basicsize */
4208 0, /* tp_itemsize */
4209 /* methods */
4210 (destructor)repeat_dealloc, /* tp_dealloc */
4211 0, /* tp_print */
4212 0, /* tp_getattr */
4213 0, /* tp_setattr */
4214 0, /* tp_reserved */
4215 (reprfunc)repeat_repr, /* tp_repr */
4216 0, /* tp_as_number */
4217 0, /* tp_as_sequence */
4218 0, /* tp_as_mapping */
4219 0, /* tp_hash */
4220 0, /* tp_call */
4221 0, /* tp_str */
4222 PyObject_GenericGetAttr, /* tp_getattro */
4223 0, /* tp_setattro */
4224 0, /* tp_as_buffer */
4225 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4226 Py_TPFLAGS_BASETYPE, /* tp_flags */
4227 repeat_doc, /* tp_doc */
4228 (traverseproc)repeat_traverse, /* tp_traverse */
4229 0, /* tp_clear */
4230 0, /* tp_richcompare */
4231 0, /* tp_weaklistoffset */
4232 PyObject_SelfIter, /* tp_iter */
4233 (iternextfunc)repeat_next, /* tp_iternext */
4234 repeat_methods, /* tp_methods */
4235 0, /* tp_members */
4236 0, /* tp_getset */
4237 0, /* tp_base */
4238 0, /* tp_dict */
4239 0, /* tp_descr_get */
4240 0, /* tp_descr_set */
4241 0, /* tp_dictoffset */
4242 0, /* tp_init */
4243 0, /* tp_alloc */
4244 repeat_new, /* tp_new */
4245 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004246};
4247
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004248/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004249
4250#include "Python.h"
4251
4252typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004253 PyObject_HEAD
4254 Py_ssize_t tuplesize;
4255 Py_ssize_t numactive;
4256 PyObject *ittuple; /* tuple of iterators */
4257 PyObject *result;
4258 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004259} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004260
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004261static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004262
4263static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004264zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004265{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004266 ziplongestobject *lz;
4267 Py_ssize_t i;
4268 PyObject *ittuple; /* tuple of iterators */
4269 PyObject *result;
4270 PyObject *fillvalue = Py_None;
4271 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004272
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004273 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4274 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4275 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4276 PyErr_SetString(PyExc_TypeError,
4277 "zip_longest() got an unexpected keyword argument");
4278 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004279 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004280 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004281
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004282 /* args must be a tuple */
4283 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004284
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004285 /* obtain iterators */
4286 ittuple = PyTuple_New(tuplesize);
4287 if (ittuple == NULL)
4288 return NULL;
4289 for (i=0; i < tuplesize; ++i) {
4290 PyObject *item = PyTuple_GET_ITEM(args, i);
4291 PyObject *it = PyObject_GetIter(item);
4292 if (it == NULL) {
4293 if (PyErr_ExceptionMatches(PyExc_TypeError))
4294 PyErr_Format(PyExc_TypeError,
4295 "zip_longest argument #%zd must support iteration",
4296 i+1);
4297 Py_DECREF(ittuple);
4298 return NULL;
4299 }
4300 PyTuple_SET_ITEM(ittuple, i, it);
4301 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004302
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004303 /* create a result holder */
4304 result = PyTuple_New(tuplesize);
4305 if (result == NULL) {
4306 Py_DECREF(ittuple);
4307 return NULL;
4308 }
4309 for (i=0 ; i < tuplesize ; i++) {
4310 Py_INCREF(Py_None);
4311 PyTuple_SET_ITEM(result, i, Py_None);
4312 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004313
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004314 /* create ziplongestobject structure */
4315 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4316 if (lz == NULL) {
4317 Py_DECREF(ittuple);
4318 Py_DECREF(result);
4319 return NULL;
4320 }
4321 lz->ittuple = ittuple;
4322 lz->tuplesize = tuplesize;
4323 lz->numactive = tuplesize;
4324 lz->result = result;
4325 Py_INCREF(fillvalue);
4326 lz->fillvalue = fillvalue;
4327 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004328}
4329
4330static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004331zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004332{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004333 PyObject_GC_UnTrack(lz);
4334 Py_XDECREF(lz->ittuple);
4335 Py_XDECREF(lz->result);
4336 Py_XDECREF(lz->fillvalue);
4337 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004338}
4339
4340static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004341zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004342{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004343 Py_VISIT(lz->ittuple);
4344 Py_VISIT(lz->result);
4345 Py_VISIT(lz->fillvalue);
4346 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004347}
4348
4349static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004350zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004351{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004352 Py_ssize_t i;
4353 Py_ssize_t tuplesize = lz->tuplesize;
4354 PyObject *result = lz->result;
4355 PyObject *it;
4356 PyObject *item;
4357 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004358
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004359 if (tuplesize == 0)
4360 return NULL;
4361 if (lz->numactive == 0)
4362 return NULL;
4363 if (Py_REFCNT(result) == 1) {
4364 Py_INCREF(result);
4365 for (i=0 ; i < tuplesize ; i++) {
4366 it = PyTuple_GET_ITEM(lz->ittuple, i);
4367 if (it == NULL) {
4368 Py_INCREF(lz->fillvalue);
4369 item = lz->fillvalue;
4370 } else {
4371 item = PyIter_Next(it);
4372 if (item == NULL) {
4373 lz->numactive -= 1;
4374 if (lz->numactive == 0 || PyErr_Occurred()) {
4375 lz->numactive = 0;
4376 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004377 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004378 } else {
4379 Py_INCREF(lz->fillvalue);
4380 item = lz->fillvalue;
4381 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4382 Py_DECREF(it);
4383 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004384 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004385 }
4386 olditem = PyTuple_GET_ITEM(result, i);
4387 PyTuple_SET_ITEM(result, i, item);
4388 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004389 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004390 } else {
4391 result = PyTuple_New(tuplesize);
4392 if (result == NULL)
4393 return NULL;
4394 for (i=0 ; i < tuplesize ; i++) {
4395 it = PyTuple_GET_ITEM(lz->ittuple, i);
4396 if (it == NULL) {
4397 Py_INCREF(lz->fillvalue);
4398 item = lz->fillvalue;
4399 } else {
4400 item = PyIter_Next(it);
4401 if (item == NULL) {
4402 lz->numactive -= 1;
4403 if (lz->numactive == 0 || PyErr_Occurred()) {
4404 lz->numactive = 0;
4405 Py_DECREF(result);
4406 return NULL;
4407 } else {
4408 Py_INCREF(lz->fillvalue);
4409 item = lz->fillvalue;
4410 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4411 Py_DECREF(it);
4412 }
4413 }
4414 }
4415 PyTuple_SET_ITEM(result, i, item);
4416 }
4417 }
4418 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004419}
4420
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004421static PyObject *
4422zip_longest_reduce(ziplongestobject *lz)
4423{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004424
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004425 /* Create a new tuple with empty sequences where appropriate to pickle.
4426 * Then use setstate to set the fillvalue
4427 */
4428 int i;
4429 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4430 if (args == NULL)
4431 return NULL;
4432 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4433 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4434 if (elem == NULL) {
4435 elem = PyTuple_New(0);
4436 if (elem == NULL) {
4437 Py_DECREF(args);
4438 return NULL;
4439 }
4440 } else
4441 Py_INCREF(elem);
4442 PyTuple_SET_ITEM(args, i, elem);
4443 }
4444 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4445}
4446
4447static PyObject *
4448zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4449{
4450 Py_CLEAR(lz->fillvalue);
4451 lz->fillvalue = state;
4452 Py_INCREF(lz->fillvalue);
4453 Py_RETURN_NONE;
4454}
4455
4456static PyMethodDef zip_longest_methods[] = {
4457 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4458 reduce_doc},
4459 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4460 setstate_doc},
4461 {NULL, NULL} /* sentinel */
4462};
4463
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004464PyDoc_STRVAR(zip_longest_doc,
4465"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004466\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004467Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004468the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004469method continues until the longest iterable in the argument sequence\n\
4470is exhausted and then it raises StopIteration. When the shorter iterables\n\
4471are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4472defaults to None or can be specified by a keyword argument.\n\
4473");
4474
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004475static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004476 PyVarObject_HEAD_INIT(NULL, 0)
4477 "itertools.zip_longest", /* tp_name */
4478 sizeof(ziplongestobject), /* tp_basicsize */
4479 0, /* tp_itemsize */
4480 /* methods */
4481 (destructor)zip_longest_dealloc, /* tp_dealloc */
4482 0, /* tp_print */
4483 0, /* tp_getattr */
4484 0, /* tp_setattr */
4485 0, /* tp_reserved */
4486 0, /* tp_repr */
4487 0, /* tp_as_number */
4488 0, /* tp_as_sequence */
4489 0, /* tp_as_mapping */
4490 0, /* tp_hash */
4491 0, /* tp_call */
4492 0, /* tp_str */
4493 PyObject_GenericGetAttr, /* tp_getattro */
4494 0, /* tp_setattro */
4495 0, /* tp_as_buffer */
4496 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4497 Py_TPFLAGS_BASETYPE, /* tp_flags */
4498 zip_longest_doc, /* tp_doc */
4499 (traverseproc)zip_longest_traverse, /* tp_traverse */
4500 0, /* tp_clear */
4501 0, /* tp_richcompare */
4502 0, /* tp_weaklistoffset */
4503 PyObject_SelfIter, /* tp_iter */
4504 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004505 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004506 0, /* tp_members */
4507 0, /* tp_getset */
4508 0, /* tp_base */
4509 0, /* tp_dict */
4510 0, /* tp_descr_get */
4511 0, /* tp_descr_set */
4512 0, /* tp_dictoffset */
4513 0, /* tp_init */
4514 0, /* tp_alloc */
4515 zip_longest_new, /* tp_new */
4516 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004517};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004518
4519/* module level code ********************************************************/
4520
4521PyDoc_STRVAR(module_doc,
4522"Functional tools for creating and using iterators.\n\
4523\n\
4524Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004525count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004526cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004527repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004528\n\
4529Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004530accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004531chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004532chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004533compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4534dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4535groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004536filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004537islice(seq, [start,] stop [, step]) --> elements from\n\
4538 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004539starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004540tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004541takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004542zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004543\n\
4544Combinatoric generators:\n\
4545product(p, q, ... [repeat=1]) --> cartesian product\n\
4546permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004547combinations(p, r)\n\
4548combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004549");
4550
4551
Raymond Hettingerad983e72003-11-12 14:32:26 +00004552static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004553 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4554 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004555};
4556
Martin v. Löwis1a214512008-06-11 05:26:20 +00004557
4558static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004559 PyModuleDef_HEAD_INIT,
4560 "itertools",
4561 module_doc,
4562 -1,
4563 module_methods,
4564 NULL,
4565 NULL,
4566 NULL,
4567 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004568};
4569
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004570PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004571PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004572{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004573 int i;
4574 PyObject *m;
4575 char *name;
4576 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004577 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004578 &combinations_type,
4579 &cwr_type,
4580 &cycle_type,
4581 &dropwhile_type,
4582 &takewhile_type,
4583 &islice_type,
4584 &starmap_type,
4585 &chain_type,
4586 &compress_type,
4587 &filterfalse_type,
4588 &count_type,
4589 &ziplongest_type,
4590 &permutations_type,
4591 &product_type,
4592 &repeat_type,
4593 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004594 &_grouper_type,
4595 &tee_type,
4596 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004597 NULL
4598 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004599
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004600 Py_TYPE(&teedataobject_type) = &PyType_Type;
4601 m = PyModule_Create(&itertoolsmodule);
4602 if (m == NULL)
4603 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004604
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004605 for (i=0 ; typelist[i] != NULL ; i++) {
4606 if (PyType_Ready(typelist[i]) < 0)
4607 return NULL;
4608 name = strchr(typelist[i]->tp_name, '.');
4609 assert (name != NULL);
4610 Py_INCREF(typelist[i]);
4611 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4612 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004613
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004614 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004615}