blob: db7cdfeefeb5404b4d86f3d7f87c0e52c10a6f3e [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 Pitrouf95a1b32010-05-09 15:52:27 +00001495 iternext = *Py_TYPE(it)->tp_iternext;
1496 while (lz->cnt < lz->next) {
1497 item = iternext(it);
1498 if (item == NULL)
1499 return NULL;
1500 Py_DECREF(item);
1501 lz->cnt++;
1502 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001503 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 return NULL;
1505 item = iternext(it);
1506 if (item == NULL)
1507 return NULL;
1508 lz->cnt++;
1509 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001510 /* The (size_t) cast below avoids the danger of undefined
1511 behaviour from signed integer overflow. */
1512 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001513 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1514 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001515 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516}
1517
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001518static PyObject *
1519islice_reduce(isliceobject *lz)
1520{
1521 /* When unpickled, generate a new object with the same bounds,
1522 * then 'setstate' with the next and count
1523 */
1524 PyObject *stop;
1525 if (lz->stop == -1) {
1526 stop = Py_None;
1527 Py_INCREF(stop);
1528 } else {
1529 stop = PyLong_FromSsize_t(lz->stop);
1530 if (stop == NULL)
1531 return NULL;
1532 }
1533 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1534 lz->it, lz->next, stop, lz->step,
1535 lz->cnt);
1536}
1537
1538static PyObject *
1539islice_setstate(isliceobject *lz, PyObject *state)
1540{
1541 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1542 if (cnt == -1 && PyErr_Occurred())
1543 return NULL;
1544 lz->cnt = cnt;
1545 Py_RETURN_NONE;
1546}
1547
1548static PyMethodDef islice_methods[] = {
1549 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1550 reduce_doc},
1551 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1552 setstate_doc},
1553 {NULL, NULL} /* sentinel */
1554};
1555
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001557"islice(iterable, stop) --> islice object\n\
1558islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559\n\
1560Return an iterator whose next() method returns selected values from an\n\
1561iterable. If start is specified, will skip all preceding elements;\n\
1562otherwise, start defaults to zero. Step defaults to one. If\n\
1563specified as another value, step determines how many values are \n\
1564skipped between successive calls. Works like a slice() on a list\n\
1565but returns an iterator.");
1566
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001567static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001568 PyVarObject_HEAD_INIT(NULL, 0)
1569 "itertools.islice", /* tp_name */
1570 sizeof(isliceobject), /* tp_basicsize */
1571 0, /* tp_itemsize */
1572 /* methods */
1573 (destructor)islice_dealloc, /* tp_dealloc */
1574 0, /* tp_print */
1575 0, /* tp_getattr */
1576 0, /* tp_setattr */
1577 0, /* tp_reserved */
1578 0, /* tp_repr */
1579 0, /* tp_as_number */
1580 0, /* tp_as_sequence */
1581 0, /* tp_as_mapping */
1582 0, /* tp_hash */
1583 0, /* tp_call */
1584 0, /* tp_str */
1585 PyObject_GenericGetAttr, /* tp_getattro */
1586 0, /* tp_setattro */
1587 0, /* tp_as_buffer */
1588 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1589 Py_TPFLAGS_BASETYPE, /* tp_flags */
1590 islice_doc, /* tp_doc */
1591 (traverseproc)islice_traverse, /* tp_traverse */
1592 0, /* tp_clear */
1593 0, /* tp_richcompare */
1594 0, /* tp_weaklistoffset */
1595 PyObject_SelfIter, /* tp_iter */
1596 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001597 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001598 0, /* tp_members */
1599 0, /* tp_getset */
1600 0, /* tp_base */
1601 0, /* tp_dict */
1602 0, /* tp_descr_get */
1603 0, /* tp_descr_set */
1604 0, /* tp_dictoffset */
1605 0, /* tp_init */
1606 0, /* tp_alloc */
1607 islice_new, /* tp_new */
1608 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609};
1610
1611
1612/* starmap object ************************************************************/
1613
1614typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001615 PyObject_HEAD
1616 PyObject *func;
1617 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618} starmapobject;
1619
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001620static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621
1622static PyObject *
1623starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1624{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001625 PyObject *func, *seq;
1626 PyObject *it;
1627 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001629 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1630 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001631
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001632 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1633 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001635 /* Get iterator. */
1636 it = PyObject_GetIter(seq);
1637 if (it == NULL)
1638 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001640 /* create starmapobject structure */
1641 lz = (starmapobject *)type->tp_alloc(type, 0);
1642 if (lz == NULL) {
1643 Py_DECREF(it);
1644 return NULL;
1645 }
1646 Py_INCREF(func);
1647 lz->func = func;
1648 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001649
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001650 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001651}
1652
1653static void
1654starmap_dealloc(starmapobject *lz)
1655{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001656 PyObject_GC_UnTrack(lz);
1657 Py_XDECREF(lz->func);
1658 Py_XDECREF(lz->it);
1659 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660}
1661
1662static int
1663starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1664{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001665 Py_VISIT(lz->it);
1666 Py_VISIT(lz->func);
1667 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668}
1669
1670static PyObject *
1671starmap_next(starmapobject *lz)
1672{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001673 PyObject *args;
1674 PyObject *result;
1675 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001677 args = (*Py_TYPE(it)->tp_iternext)(it);
1678 if (args == NULL)
1679 return NULL;
1680 if (!PyTuple_CheckExact(args)) {
1681 PyObject *newargs = PySequence_Tuple(args);
1682 Py_DECREF(args);
1683 if (newargs == NULL)
1684 return NULL;
1685 args = newargs;
1686 }
1687 result = PyObject_Call(lz->func, args, NULL);
1688 Py_DECREF(args);
1689 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690}
1691
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001692static PyObject *
1693starmap_reduce(starmapobject *lz)
1694{
1695 /* Just pickle the iterator */
1696 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1697}
1698
1699static PyMethodDef starmap_methods[] = {
1700 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1701 reduce_doc},
1702 {NULL, NULL} /* sentinel */
1703};
1704
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705PyDoc_STRVAR(starmap_doc,
1706"starmap(function, sequence) --> starmap object\n\
1707\n\
1708Return an iterator whose values are returned from the function evaluated\n\
1709with a argument tuple taken from the given sequence.");
1710
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001711static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001712 PyVarObject_HEAD_INIT(NULL, 0)
1713 "itertools.starmap", /* tp_name */
1714 sizeof(starmapobject), /* tp_basicsize */
1715 0, /* tp_itemsize */
1716 /* methods */
1717 (destructor)starmap_dealloc, /* tp_dealloc */
1718 0, /* tp_print */
1719 0, /* tp_getattr */
1720 0, /* tp_setattr */
1721 0, /* tp_reserved */
1722 0, /* tp_repr */
1723 0, /* tp_as_number */
1724 0, /* tp_as_sequence */
1725 0, /* tp_as_mapping */
1726 0, /* tp_hash */
1727 0, /* tp_call */
1728 0, /* tp_str */
1729 PyObject_GenericGetAttr, /* tp_getattro */
1730 0, /* tp_setattro */
1731 0, /* tp_as_buffer */
1732 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1733 Py_TPFLAGS_BASETYPE, /* tp_flags */
1734 starmap_doc, /* tp_doc */
1735 (traverseproc)starmap_traverse, /* tp_traverse */
1736 0, /* tp_clear */
1737 0, /* tp_richcompare */
1738 0, /* tp_weaklistoffset */
1739 PyObject_SelfIter, /* tp_iter */
1740 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001741 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001742 0, /* tp_members */
1743 0, /* tp_getset */
1744 0, /* tp_base */
1745 0, /* tp_dict */
1746 0, /* tp_descr_get */
1747 0, /* tp_descr_set */
1748 0, /* tp_dictoffset */
1749 0, /* tp_init */
1750 0, /* tp_alloc */
1751 starmap_new, /* tp_new */
1752 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753};
1754
1755
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001756/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757
1758typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001759 PyObject_HEAD
1760 PyObject *source; /* Iterator over input iterables */
1761 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001762} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001763
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001764static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001765
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001766static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001767chain_new_internal(PyTypeObject *type, PyObject *source)
1768{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001769 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001771 lz = (chainobject *)type->tp_alloc(type, 0);
1772 if (lz == NULL) {
1773 Py_DECREF(source);
1774 return NULL;
1775 }
1776
1777 lz->source = source;
1778 lz->active = NULL;
1779 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001780}
1781
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001782static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001783chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001786
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001787 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1788 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001789
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001790 source = PyObject_GetIter(args);
1791 if (source == NULL)
1792 return NULL;
1793
1794 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001795}
1796
1797static PyObject *
1798chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1799{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001800 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001801
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001802 source = PyObject_GetIter(arg);
1803 if (source == NULL)
1804 return NULL;
1805
1806 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001807}
1808
1809static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001810chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001811{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001812 PyObject_GC_UnTrack(lz);
1813 Py_XDECREF(lz->active);
1814 Py_XDECREF(lz->source);
1815 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001816}
1817
Raymond Hettinger2012f172003-02-07 05:32:58 +00001818static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001819chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001820{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 Py_VISIT(lz->source);
1822 Py_VISIT(lz->active);
1823 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001824}
1825
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001827chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001828{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001829 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001831 if (lz->source == NULL)
1832 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001833
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001834 if (lz->active == NULL) {
1835 PyObject *iterable = PyIter_Next(lz->source);
1836 if (iterable == NULL) {
1837 Py_CLEAR(lz->source);
1838 return NULL; /* no more input sources */
1839 }
1840 lz->active = PyObject_GetIter(iterable);
1841 Py_DECREF(iterable);
1842 if (lz->active == NULL) {
1843 Py_CLEAR(lz->source);
1844 return NULL; /* input not iterable */
1845 }
1846 }
1847 item = PyIter_Next(lz->active);
1848 if (item != NULL)
1849 return item;
1850 if (PyErr_Occurred()) {
1851 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1852 PyErr_Clear();
1853 else
1854 return NULL; /* input raised an exception */
1855 }
1856 Py_CLEAR(lz->active);
1857 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001858}
1859
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001860static PyObject *
1861chain_reduce(chainobject *lz)
1862{
1863 if (lz->source) {
1864 /* we can't pickle function objects (itertools.from_iterable) so
1865 * we must use setstate to replace the iterable. One day we
1866 * will fix pickling of functions
1867 */
1868 if (lz->active) {
1869 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1870 } else {
1871 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1872 }
1873 } else {
1874 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1875 }
1876 return NULL;
1877}
1878
1879static PyObject *
1880chain_setstate(chainobject *lz, PyObject *state)
1881{
1882 PyObject *source, *active=NULL;
1883 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1884 return NULL;
1885
1886 Py_CLEAR(lz->source);
1887 lz->source = source;
1888 Py_INCREF(lz->source);
1889 Py_CLEAR(lz->active);
1890 lz->active = active;
1891 Py_XINCREF(lz->active);
1892 Py_RETURN_NONE;
1893}
1894
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001895PyDoc_STRVAR(chain_doc,
1896"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001897\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001898Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001899first iterable until it is exhausted, then elements from the next\n\
1900iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001901
Christian Heimesf16baeb2008-02-29 14:57:44 +00001902PyDoc_STRVAR(chain_from_iterable_doc,
1903"chain.from_iterable(iterable) --> chain object\n\
1904\n\
1905Alternate chain() contructor taking a single iterable argument\n\
1906that evaluates lazily.");
1907
1908static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001909 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1910 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001911 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1912 reduce_doc},
1913 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1914 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001915 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001916};
1917
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001918static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001919 PyVarObject_HEAD_INIT(NULL, 0)
1920 "itertools.chain", /* tp_name */
1921 sizeof(chainobject), /* tp_basicsize */
1922 0, /* tp_itemsize */
1923 /* methods */
1924 (destructor)chain_dealloc, /* tp_dealloc */
1925 0, /* tp_print */
1926 0, /* tp_getattr */
1927 0, /* tp_setattr */
1928 0, /* tp_reserved */
1929 0, /* tp_repr */
1930 0, /* tp_as_number */
1931 0, /* tp_as_sequence */
1932 0, /* tp_as_mapping */
1933 0, /* tp_hash */
1934 0, /* tp_call */
1935 0, /* tp_str */
1936 PyObject_GenericGetAttr, /* tp_getattro */
1937 0, /* tp_setattro */
1938 0, /* tp_as_buffer */
1939 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1940 Py_TPFLAGS_BASETYPE, /* tp_flags */
1941 chain_doc, /* tp_doc */
1942 (traverseproc)chain_traverse, /* tp_traverse */
1943 0, /* tp_clear */
1944 0, /* tp_richcompare */
1945 0, /* tp_weaklistoffset */
1946 PyObject_SelfIter, /* tp_iter */
1947 (iternextfunc)chain_next, /* tp_iternext */
1948 chain_methods, /* tp_methods */
1949 0, /* tp_members */
1950 0, /* tp_getset */
1951 0, /* tp_base */
1952 0, /* tp_dict */
1953 0, /* tp_descr_get */
1954 0, /* tp_descr_set */
1955 0, /* tp_dictoffset */
1956 0, /* tp_init */
1957 0, /* tp_alloc */
1958 chain_new, /* tp_new */
1959 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001960};
1961
1962
Christian Heimesc3f30c42008-02-22 16:37:40 +00001963/* product object ************************************************************/
1964
1965typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001966 PyObject_HEAD
1967 PyObject *pools; /* tuple of pool tuples */
1968 Py_ssize_t *indices; /* one index per pool */
1969 PyObject *result; /* most recently returned result tuple */
1970 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001971} productobject;
1972
1973static PyTypeObject product_type;
1974
1975static PyObject *
1976product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1977{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001978 productobject *lz;
1979 Py_ssize_t nargs, npools, repeat=1;
1980 PyObject *pools = NULL;
1981 Py_ssize_t *indices = NULL;
1982 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001983
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001984 if (kwds != NULL) {
1985 char *kwlist[] = {"repeat", 0};
1986 PyObject *tmpargs = PyTuple_New(0);
1987 if (tmpargs == NULL)
1988 return NULL;
1989 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1990 Py_DECREF(tmpargs);
1991 return NULL;
1992 }
1993 Py_DECREF(tmpargs);
1994 if (repeat < 0) {
1995 PyErr_SetString(PyExc_ValueError,
1996 "repeat argument cannot be negative");
1997 return NULL;
1998 }
1999 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002001 assert(PyTuple_Check(args));
2002 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
2003 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002005 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
2006 if (indices == NULL) {
2007 PyErr_NoMemory();
2008 goto error;
2009 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002010
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002011 pools = PyTuple_New(npools);
2012 if (pools == NULL)
2013 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002014
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002015 for (i=0; i < nargs ; ++i) {
2016 PyObject *item = PyTuple_GET_ITEM(args, i);
2017 PyObject *pool = PySequence_Tuple(item);
2018 if (pool == NULL)
2019 goto error;
2020 PyTuple_SET_ITEM(pools, i, pool);
2021 indices[i] = 0;
2022 }
2023 for ( ; i < npools; ++i) {
2024 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2025 Py_INCREF(pool);
2026 PyTuple_SET_ITEM(pools, i, pool);
2027 indices[i] = 0;
2028 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002029
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002030 /* create productobject structure */
2031 lz = (productobject *)type->tp_alloc(type, 0);
2032 if (lz == NULL)
2033 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002035 lz->pools = pools;
2036 lz->indices = indices;
2037 lz->result = NULL;
2038 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002039
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002040 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002041
2042error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002043 if (indices != NULL)
2044 PyMem_Free(indices);
2045 Py_XDECREF(pools);
2046 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002047}
2048
2049static void
2050product_dealloc(productobject *lz)
2051{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002052 PyObject_GC_UnTrack(lz);
2053 Py_XDECREF(lz->pools);
2054 Py_XDECREF(lz->result);
2055 if (lz->indices != NULL)
2056 PyMem_Free(lz->indices);
2057 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002058}
2059
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002060static PyObject *
2061product_sizeof(productobject *lz, void *unused)
2062{
2063 Py_ssize_t res;
2064
2065 res = sizeof(productobject);
2066 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2067 return PyLong_FromSsize_t(res);
2068}
2069
2070PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2071
Christian Heimesc3f30c42008-02-22 16:37:40 +00002072static int
2073product_traverse(productobject *lz, visitproc visit, void *arg)
2074{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002075 Py_VISIT(lz->pools);
2076 Py_VISIT(lz->result);
2077 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002078}
2079
2080static PyObject *
2081product_next(productobject *lz)
2082{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002083 PyObject *pool;
2084 PyObject *elem;
2085 PyObject *oldelem;
2086 PyObject *pools = lz->pools;
2087 PyObject *result = lz->result;
2088 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2089 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002090
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002091 if (lz->stopped)
2092 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002093
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002094 if (result == NULL) {
2095 /* On the first pass, return an initial tuple filled with the
2096 first element from each pool. */
2097 result = PyTuple_New(npools);
2098 if (result == NULL)
2099 goto empty;
2100 lz->result = result;
2101 for (i=0; i < npools; i++) {
2102 pool = PyTuple_GET_ITEM(pools, i);
2103 if (PyTuple_GET_SIZE(pool) == 0)
2104 goto empty;
2105 elem = PyTuple_GET_ITEM(pool, 0);
2106 Py_INCREF(elem);
2107 PyTuple_SET_ITEM(result, i, elem);
2108 }
2109 } else {
2110 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002112 /* Copy the previous result tuple or re-use it if available */
2113 if (Py_REFCNT(result) > 1) {
2114 PyObject *old_result = result;
2115 result = PyTuple_New(npools);
2116 if (result == NULL)
2117 goto empty;
2118 lz->result = result;
2119 for (i=0; i < npools; i++) {
2120 elem = PyTuple_GET_ITEM(old_result, i);
2121 Py_INCREF(elem);
2122 PyTuple_SET_ITEM(result, i, elem);
2123 }
2124 Py_DECREF(old_result);
2125 }
2126 /* Now, we've got the only copy so we can update it in-place */
2127 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002128
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002129 /* Update the pool indices right-to-left. Only advance to the
2130 next pool when the previous one rolls-over */
2131 for (i=npools-1 ; i >= 0 ; i--) {
2132 pool = PyTuple_GET_ITEM(pools, i);
2133 indices[i]++;
2134 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2135 /* Roll-over and advance to next pool */
2136 indices[i] = 0;
2137 elem = PyTuple_GET_ITEM(pool, 0);
2138 Py_INCREF(elem);
2139 oldelem = PyTuple_GET_ITEM(result, i);
2140 PyTuple_SET_ITEM(result, i, elem);
2141 Py_DECREF(oldelem);
2142 } else {
2143 /* No rollover. Just increment and stop here. */
2144 elem = PyTuple_GET_ITEM(pool, indices[i]);
2145 Py_INCREF(elem);
2146 oldelem = PyTuple_GET_ITEM(result, i);
2147 PyTuple_SET_ITEM(result, i, elem);
2148 Py_DECREF(oldelem);
2149 break;
2150 }
2151 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002152
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002153 /* If i is negative, then the indices have all rolled-over
2154 and we're done. */
2155 if (i < 0)
2156 goto empty;
2157 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002159 Py_INCREF(result);
2160 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002161
2162empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002163 lz->stopped = 1;
2164 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002165}
2166
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002167static PyObject *
2168product_reduce(productobject *lz)
2169{
2170 if (lz->stopped) {
2171 return Py_BuildValue("O(())", Py_TYPE(lz));
2172 } else if (lz->result == NULL) {
2173 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2174 } else {
2175 PyObject *indices;
2176 Py_ssize_t n, i;
2177
2178 /* we must pickle the indices use them for setstate, and
2179 * additionally indicate that the iterator has started
2180 */
2181 n = PyTuple_GET_SIZE(lz->pools);
2182 indices = PyTuple_New(n);
2183 if (indices == NULL)
2184 return NULL;
2185 for (i=0; i<n; i++){
2186 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2187 if (!index) {
2188 Py_DECREF(indices);
2189 return NULL;
2190 }
2191 PyTuple_SET_ITEM(indices, i, index);
2192 }
2193 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2194 }
2195}
2196
2197static PyObject *
2198product_setstate(productobject *lz, PyObject *state)
2199{
2200 PyObject *result;
2201 Py_ssize_t n, i;
2202
2203 n = PyTuple_GET_SIZE(lz->pools);
2204 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2205 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2206 return NULL;
2207 }
2208 for (i=0; i<n; i++)
2209 {
2210 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2211 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2212 if (index < 0 && PyErr_Occurred())
2213 return NULL; /* not an integer */
2214 /* clamp the index */
2215 if (index < 0)
2216 index = 0;
2217 else if (index > n-1)
2218 index = n-1;
2219 lz->indices[i] = index;
2220 }
2221
2222 result = PyTuple_New(n);
2223 if (!result)
2224 return NULL;
2225 for (i=0; i<n; i++) {
2226 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2227 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2228 Py_INCREF(element);
2229 PyTuple_SET_ITEM(result, i, element);
2230 }
2231 Py_CLEAR(lz->result);
2232 lz->result = result;
2233 Py_RETURN_NONE;
2234}
2235
2236static PyMethodDef product_methods[] = {
2237 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2238 reduce_doc},
2239 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2240 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002241 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2242 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002243 {NULL, NULL} /* sentinel */
2244};
2245
Christian Heimesc3f30c42008-02-22 16:37:40 +00002246PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002247"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002248\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002249Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002250For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2251The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2252cycle in a manner similar to an odometer (with the rightmost element changing\n\
2253on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002254To compute the product of an iterable with itself, specify the number\n\
2255of repetitions with the optional repeat keyword argument. For example,\n\
2256product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002257product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2258product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2259
2260static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002261 PyVarObject_HEAD_INIT(NULL, 0)
2262 "itertools.product", /* tp_name */
2263 sizeof(productobject), /* tp_basicsize */
2264 0, /* tp_itemsize */
2265 /* methods */
2266 (destructor)product_dealloc, /* tp_dealloc */
2267 0, /* tp_print */
2268 0, /* tp_getattr */
2269 0, /* tp_setattr */
2270 0, /* tp_reserved */
2271 0, /* tp_repr */
2272 0, /* tp_as_number */
2273 0, /* tp_as_sequence */
2274 0, /* tp_as_mapping */
2275 0, /* tp_hash */
2276 0, /* tp_call */
2277 0, /* tp_str */
2278 PyObject_GenericGetAttr, /* tp_getattro */
2279 0, /* tp_setattro */
2280 0, /* tp_as_buffer */
2281 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2282 Py_TPFLAGS_BASETYPE, /* tp_flags */
2283 product_doc, /* tp_doc */
2284 (traverseproc)product_traverse, /* tp_traverse */
2285 0, /* tp_clear */
2286 0, /* tp_richcompare */
2287 0, /* tp_weaklistoffset */
2288 PyObject_SelfIter, /* tp_iter */
2289 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002290 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002291 0, /* tp_members */
2292 0, /* tp_getset */
2293 0, /* tp_base */
2294 0, /* tp_dict */
2295 0, /* tp_descr_get */
2296 0, /* tp_descr_set */
2297 0, /* tp_dictoffset */
2298 0, /* tp_init */
2299 0, /* tp_alloc */
2300 product_new, /* tp_new */
2301 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002302};
2303
2304
Christian Heimes380f7f22008-02-28 11:19:05 +00002305/* combinations object ************************************************************/
2306
2307typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002308 PyObject_HEAD
2309 PyObject *pool; /* input converted to a tuple */
2310 Py_ssize_t *indices; /* one index per result element */
2311 PyObject *result; /* most recently returned result tuple */
2312 Py_ssize_t r; /* size of result tuple */
2313 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002314} combinationsobject;
2315
2316static PyTypeObject combinations_type;
2317
2318static PyObject *
2319combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2320{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002321 combinationsobject *co;
2322 Py_ssize_t n;
2323 Py_ssize_t r;
2324 PyObject *pool = NULL;
2325 PyObject *iterable = NULL;
2326 Py_ssize_t *indices = NULL;
2327 Py_ssize_t i;
2328 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002329
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002330 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2331 &iterable, &r))
2332 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002333
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002334 pool = PySequence_Tuple(iterable);
2335 if (pool == NULL)
2336 goto error;
2337 n = PyTuple_GET_SIZE(pool);
2338 if (r < 0) {
2339 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2340 goto error;
2341 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002342
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002343 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2344 if (indices == NULL) {
2345 PyErr_NoMemory();
2346 goto error;
2347 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002348
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002349 for (i=0 ; i<r ; i++)
2350 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002351
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002352 /* create combinationsobject structure */
2353 co = (combinationsobject *)type->tp_alloc(type, 0);
2354 if (co == NULL)
2355 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002356
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002357 co->pool = pool;
2358 co->indices = indices;
2359 co->result = NULL;
2360 co->r = r;
2361 co->stopped = r > n ? 1 : 0;
2362
2363 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002364
2365error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002366 if (indices != NULL)
2367 PyMem_Free(indices);
2368 Py_XDECREF(pool);
2369 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002370}
2371
2372static void
2373combinations_dealloc(combinationsobject *co)
2374{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002375 PyObject_GC_UnTrack(co);
2376 Py_XDECREF(co->pool);
2377 Py_XDECREF(co->result);
2378 if (co->indices != NULL)
2379 PyMem_Free(co->indices);
2380 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002381}
2382
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002383static PyObject *
2384combinations_sizeof(combinationsobject *co, void *unused)
2385{
2386 Py_ssize_t res;
2387
2388 res = sizeof(combinationsobject);
2389 res += co->r * sizeof(Py_ssize_t);
2390 return PyLong_FromSsize_t(res);
2391}
2392
Christian Heimes380f7f22008-02-28 11:19:05 +00002393static int
2394combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2395{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002396 Py_VISIT(co->pool);
2397 Py_VISIT(co->result);
2398 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002399}
2400
2401static PyObject *
2402combinations_next(combinationsobject *co)
2403{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002404 PyObject *elem;
2405 PyObject *oldelem;
2406 PyObject *pool = co->pool;
2407 Py_ssize_t *indices = co->indices;
2408 PyObject *result = co->result;
2409 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2410 Py_ssize_t r = co->r;
2411 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002412
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002413 if (co->stopped)
2414 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002416 if (result == NULL) {
2417 /* On the first pass, initialize result tuple using the indices */
2418 result = PyTuple_New(r);
2419 if (result == NULL)
2420 goto empty;
2421 co->result = result;
2422 for (i=0; i<r ; i++) {
2423 index = indices[i];
2424 elem = PyTuple_GET_ITEM(pool, index);
2425 Py_INCREF(elem);
2426 PyTuple_SET_ITEM(result, i, elem);
2427 }
2428 } else {
2429 /* Copy the previous result tuple or re-use it if available */
2430 if (Py_REFCNT(result) > 1) {
2431 PyObject *old_result = result;
2432 result = PyTuple_New(r);
2433 if (result == NULL)
2434 goto empty;
2435 co->result = result;
2436 for (i=0; i<r ; i++) {
2437 elem = PyTuple_GET_ITEM(old_result, i);
2438 Py_INCREF(elem);
2439 PyTuple_SET_ITEM(result, i, elem);
2440 }
2441 Py_DECREF(old_result);
2442 }
2443 /* Now, we've got the only copy so we can update it in-place
2444 * CPython's empty tuple is a singleton and cached in
2445 * PyTuple's freelist.
2446 */
2447 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002449 /* Scan indices right-to-left until finding one that is not
2450 at its maximum (i + n - r). */
2451 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2452 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002453
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002454 /* If i is negative, then the indices are all at
2455 their maximum value and we're done. */
2456 if (i < 0)
2457 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002458
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002459 /* Increment the current index which we know is not at its
2460 maximum. Then move back to the right setting each index
2461 to its lowest possible value (one higher than the index
2462 to its left -- this maintains the sort order invariant). */
2463 indices[i]++;
2464 for (j=i+1 ; j<r ; j++)
2465 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002467 /* Update the result tuple for the new indices
2468 starting with i, the leftmost index that changed */
2469 for ( ; i<r ; i++) {
2470 index = indices[i];
2471 elem = PyTuple_GET_ITEM(pool, index);
2472 Py_INCREF(elem);
2473 oldelem = PyTuple_GET_ITEM(result, i);
2474 PyTuple_SET_ITEM(result, i, elem);
2475 Py_DECREF(oldelem);
2476 }
2477 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002478
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002479 Py_INCREF(result);
2480 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002481
2482empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002483 co->stopped = 1;
2484 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002485}
2486
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002487static PyObject *
2488combinations_reduce(combinationsobject *lz)
2489{
2490 if (lz->result == NULL) {
2491 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2492 } else if (lz->stopped) {
2493 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2494 } else {
2495 PyObject *indices;
2496 Py_ssize_t i;
2497
2498 /* we must pickle the indices and use them for setstate */
2499 indices = PyTuple_New(lz->r);
2500 if (!indices)
2501 return NULL;
2502 for (i=0; i<lz->r; i++)
2503 {
2504 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2505 if (!index) {
2506 Py_DECREF(indices);
2507 return NULL;
2508 }
2509 PyTuple_SET_ITEM(indices, i, index);
2510 }
2511
2512 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2513 }
2514}
2515
2516static PyObject *
2517combinations_setstate(combinationsobject *lz, PyObject *state)
2518{
2519 PyObject *result;
2520 Py_ssize_t i;
2521 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2522
2523 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2524 {
2525 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2526 return NULL;
2527 }
2528
2529 for (i=0; i<lz->r; i++)
2530 {
2531 Py_ssize_t max;
2532 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2533 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2534 if (index == -1 && PyErr_Occurred())
2535 return NULL; /* not an integer */
2536 max = i + n - lz->r;
2537 /* clamp the index (beware of negative max) */
2538 if (index > max)
2539 index = max;
2540 if (index < 0)
2541 index = 0;
2542 lz->indices[i] = index;
2543 }
2544
2545 result = PyTuple_New(lz->r);
2546 if (result == NULL)
2547 return NULL;
2548 for (i=0; i<lz->r; i++) {
2549 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2550 Py_INCREF(element);
2551 PyTuple_SET_ITEM(result, i, element);
2552 }
2553
2554 Py_CLEAR(lz->result);
2555 lz->result = result;
2556 Py_RETURN_NONE;
2557}
2558
2559static PyMethodDef combinations_methods[] = {
2560 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2561 reduce_doc},
2562 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2563 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002564 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2565 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002566 {NULL, NULL} /* sentinel */
2567};
2568
Christian Heimes380f7f22008-02-28 11:19:05 +00002569PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002570"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002571\n\
2572Return successive r-length combinations of elements in the iterable.\n\n\
2573combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2574
2575static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002576 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002577 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002578 sizeof(combinationsobject), /* tp_basicsize */
2579 0, /* tp_itemsize */
2580 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002581 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002582 0, /* tp_print */
2583 0, /* tp_getattr */
2584 0, /* tp_setattr */
2585 0, /* tp_reserved */
2586 0, /* tp_repr */
2587 0, /* tp_as_number */
2588 0, /* tp_as_sequence */
2589 0, /* tp_as_mapping */
2590 0, /* tp_hash */
2591 0, /* tp_call */
2592 0, /* tp_str */
2593 PyObject_GenericGetAttr, /* tp_getattro */
2594 0, /* tp_setattro */
2595 0, /* tp_as_buffer */
2596 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2597 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002598 combinations_doc, /* tp_doc */
2599 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002600 0, /* tp_clear */
2601 0, /* tp_richcompare */
2602 0, /* tp_weaklistoffset */
2603 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002604 (iternextfunc)combinations_next, /* tp_iternext */
2605 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002606 0, /* tp_members */
2607 0, /* tp_getset */
2608 0, /* tp_base */
2609 0, /* tp_dict */
2610 0, /* tp_descr_get */
2611 0, /* tp_descr_set */
2612 0, /* tp_dictoffset */
2613 0, /* tp_init */
2614 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002615 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002616 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002617};
2618
2619
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002620/* combinations with replacement object *******************************************/
2621
2622/* Equivalent to:
2623
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002624 def combinations_with_replacement(iterable, r):
2625 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2626 # number items returned: (n+r-1)! / r! / (n-1)!
2627 pool = tuple(iterable)
2628 n = len(pool)
2629 indices = [0] * r
2630 yield tuple(pool[i] for i in indices)
2631 while 1:
2632 for i in reversed(range(r)):
2633 if indices[i] != n - 1:
2634 break
2635 else:
2636 return
2637 indices[i:] = [indices[i] + 1] * (r - i)
2638 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002639
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002640 def combinations_with_replacement2(iterable, r):
2641 'Alternate version that filters from product()'
2642 pool = tuple(iterable)
2643 n = len(pool)
2644 for indices in product(range(n), repeat=r):
2645 if sorted(indices) == list(indices):
2646 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002647*/
2648typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002649 PyObject_HEAD
2650 PyObject *pool; /* input converted to a tuple */
2651 Py_ssize_t *indices; /* one index per result element */
2652 PyObject *result; /* most recently returned result tuple */
2653 Py_ssize_t r; /* size of result tuple */
2654 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002655} cwrobject;
2656
2657static PyTypeObject cwr_type;
2658
2659static PyObject *
2660cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2661{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002662 cwrobject *co;
2663 Py_ssize_t n;
2664 Py_ssize_t r;
2665 PyObject *pool = NULL;
2666 PyObject *iterable = NULL;
2667 Py_ssize_t *indices = NULL;
2668 Py_ssize_t i;
2669 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002670
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002671 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2672 &iterable, &r))
2673 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002675 pool = PySequence_Tuple(iterable);
2676 if (pool == NULL)
2677 goto error;
2678 n = PyTuple_GET_SIZE(pool);
2679 if (r < 0) {
2680 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2681 goto error;
2682 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002683
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002684 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2685 if (indices == NULL) {
2686 PyErr_NoMemory();
2687 goto error;
2688 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002689
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002690 for (i=0 ; i<r ; i++)
2691 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002692
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002693 /* create cwrobject structure */
2694 co = (cwrobject *)type->tp_alloc(type, 0);
2695 if (co == NULL)
2696 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 co->pool = pool;
2699 co->indices = indices;
2700 co->result = NULL;
2701 co->r = r;
2702 co->stopped = !n && r;
2703
2704 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002705
2706error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002707 if (indices != NULL)
2708 PyMem_Free(indices);
2709 Py_XDECREF(pool);
2710 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002711}
2712
2713static void
2714cwr_dealloc(cwrobject *co)
2715{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002716 PyObject_GC_UnTrack(co);
2717 Py_XDECREF(co->pool);
2718 Py_XDECREF(co->result);
2719 if (co->indices != NULL)
2720 PyMem_Free(co->indices);
2721 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002722}
2723
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002724static PyObject *
2725cwr_sizeof(cwrobject *co, void *unused)
2726{
2727 Py_ssize_t res;
2728
2729 res = sizeof(cwrobject);
2730 res += co->r * sizeof(Py_ssize_t);
2731 return PyLong_FromSsize_t(res);
2732}
2733
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002734static int
2735cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2736{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002737 Py_VISIT(co->pool);
2738 Py_VISIT(co->result);
2739 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002740}
2741
2742static PyObject *
2743cwr_next(cwrobject *co)
2744{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002745 PyObject *elem;
2746 PyObject *oldelem;
2747 PyObject *pool = co->pool;
2748 Py_ssize_t *indices = co->indices;
2749 PyObject *result = co->result;
2750 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2751 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002752 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002753
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002754 if (co->stopped)
2755 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002757 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002758 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002759 result = PyTuple_New(r);
2760 if (result == NULL)
2761 goto empty;
2762 co->result = result;
Tim Peters9edb1682013-09-03 11:49:31 -05002763 elem = PyTuple_GET_ITEM(pool, 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002764 for (i=0; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002765 assert(indices[i] == 0);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002766 Py_INCREF(elem);
2767 PyTuple_SET_ITEM(result, i, elem);
2768 }
2769 } else {
2770 /* Copy the previous result tuple or re-use it if available */
2771 if (Py_REFCNT(result) > 1) {
2772 PyObject *old_result = result;
2773 result = PyTuple_New(r);
2774 if (result == NULL)
2775 goto empty;
2776 co->result = result;
2777 for (i=0; i<r ; i++) {
2778 elem = PyTuple_GET_ITEM(old_result, i);
2779 Py_INCREF(elem);
2780 PyTuple_SET_ITEM(result, i, elem);
2781 }
2782 Py_DECREF(old_result);
2783 }
2784 /* Now, we've got the only copy so we can update it in-place CPython's
2785 empty tuple is a singleton and cached in PyTuple's freelist. */
2786 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002787
Tim Peters9edb1682013-09-03 11:49:31 -05002788 /* Scan indices right-to-left until finding one that is not
2789 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002790 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2791 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002792
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002794 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002795 if (i < 0)
2796 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002797
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002798 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002799 maximum. Then set all to the right to the same value. */
2800 index = indices[i] + 1;
2801 assert(index < n);
2802 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002803 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002804 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002805 Py_INCREF(elem);
2806 oldelem = PyTuple_GET_ITEM(result, i);
2807 PyTuple_SET_ITEM(result, i, elem);
2808 Py_DECREF(oldelem);
2809 }
2810 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 Py_INCREF(result);
2813 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002814
2815empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002816 co->stopped = 1;
2817 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002818}
2819
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002820static PyObject *
2821cwr_reduce(cwrobject *lz)
2822{
2823 if (lz->result == NULL) {
2824 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2825 } else if (lz->stopped) {
2826 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2827 } else {
2828 PyObject *indices;
2829 Py_ssize_t i;
2830
2831 /* we must pickle the indices and use them for setstate */
2832 indices = PyTuple_New(lz->r);
2833 if (!indices)
2834 return NULL;
2835 for (i=0; i<lz->r; i++)
2836 {
2837 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2838 if (!index) {
2839 Py_DECREF(indices);
2840 return NULL;
2841 }
2842 PyTuple_SET_ITEM(indices, i, index);
2843 }
2844
2845 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2846 }
2847}
2848
2849static PyObject *
2850cwr_setstate(cwrobject *lz, PyObject *state)
2851{
2852 PyObject *result;
2853 Py_ssize_t n, i;
2854
2855 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2856 {
2857 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2858 return NULL;
2859 }
2860
2861 n = PyTuple_GET_SIZE(lz->pool);
2862 for (i=0; i<lz->r; i++)
2863 {
2864 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2865 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2866 if (index < 0 && PyErr_Occurred())
2867 return NULL; /* not an integer */
2868 /* clamp the index */
2869 if (index < 0)
2870 index = 0;
2871 else if (index > n-1)
2872 index = n-1;
2873 lz->indices[i] = index;
2874 }
2875 result = PyTuple_New(lz->r);
2876 if (result == NULL)
2877 return NULL;
2878 for (i=0; i<lz->r; i++) {
2879 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2880 Py_INCREF(element);
2881 PyTuple_SET_ITEM(result, i, element);
2882 }
2883 Py_CLEAR(lz->result);
2884 lz->result = result;
2885 Py_RETURN_NONE;
2886}
2887
2888static PyMethodDef cwr_methods[] = {
2889 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2890 reduce_doc},
2891 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2892 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002893 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
2894 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002895 {NULL, NULL} /* sentinel */
2896};
2897
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002898PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002899"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002900\n\
2901Return successive r-length combinations of elements in the iterable\n\
2902allowing individual elements to have successive repeats.\n\
2903combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2904
2905static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002906 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002907 "itertools.combinations_with_replacement", /* tp_name */
2908 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002909 0, /* tp_itemsize */
2910 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002911 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002912 0, /* tp_print */
2913 0, /* tp_getattr */
2914 0, /* tp_setattr */
2915 0, /* tp_reserved */
2916 0, /* tp_repr */
2917 0, /* tp_as_number */
2918 0, /* tp_as_sequence */
2919 0, /* tp_as_mapping */
2920 0, /* tp_hash */
2921 0, /* tp_call */
2922 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002923 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002924 0, /* tp_setattro */
2925 0, /* tp_as_buffer */
2926 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002927 Py_TPFLAGS_BASETYPE, /* tp_flags */
2928 cwr_doc, /* tp_doc */
2929 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002930 0, /* tp_clear */
2931 0, /* tp_richcompare */
2932 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002933 PyObject_SelfIter, /* tp_iter */
2934 (iternextfunc)cwr_next, /* tp_iternext */
2935 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002936 0, /* tp_members */
2937 0, /* tp_getset */
2938 0, /* tp_base */
2939 0, /* tp_dict */
2940 0, /* tp_descr_get */
2941 0, /* tp_descr_set */
2942 0, /* tp_dictoffset */
2943 0, /* tp_init */
2944 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002945 cwr_new, /* tp_new */
2946 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002947};
2948
2949
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002950/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002951
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002952def permutations(iterable, r=None):
2953 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2954 pool = tuple(iterable)
2955 n = len(pool)
2956 r = n if r is None else r
2957 indices = range(n)
2958 cycles = range(n-r+1, n+1)[::-1]
2959 yield tuple(pool[i] for i in indices[:r])
2960 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002961 for i in reversed(range(r)):
2962 cycles[i] -= 1
2963 if cycles[i] == 0:
2964 indices[i:] = indices[i+1:] + indices[i:i+1]
2965 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002966 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002967 j = cycles[i]
2968 indices[i], indices[-j] = indices[-j], indices[i]
2969 yield tuple(pool[i] for i in indices[:r])
2970 break
2971 else:
2972 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002973*/
2974
2975typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002976 PyObject_HEAD
2977 PyObject *pool; /* input converted to a tuple */
2978 Py_ssize_t *indices; /* one index per element in the pool */
2979 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2980 PyObject *result; /* most recently returned result tuple */
2981 Py_ssize_t r; /* size of result tuple */
2982 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002983} permutationsobject;
2984
2985static PyTypeObject permutations_type;
2986
2987static PyObject *
2988permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2989{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002990 permutationsobject *po;
2991 Py_ssize_t n;
2992 Py_ssize_t r;
2993 PyObject *robj = Py_None;
2994 PyObject *pool = NULL;
2995 PyObject *iterable = NULL;
2996 Py_ssize_t *indices = NULL;
2997 Py_ssize_t *cycles = NULL;
2998 Py_ssize_t i;
2999 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003001 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
3002 &iterable, &robj))
3003 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003005 pool = PySequence_Tuple(iterable);
3006 if (pool == NULL)
3007 goto error;
3008 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003010 r = n;
3011 if (robj != Py_None) {
3012 if (!PyLong_Check(robj)) {
3013 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3014 goto error;
3015 }
3016 r = PyLong_AsSsize_t(robj);
3017 if (r == -1 && PyErr_Occurred())
3018 goto error;
3019 }
3020 if (r < 0) {
3021 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3022 goto error;
3023 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003024
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
3026 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
3027 if (indices == NULL || cycles == NULL) {
3028 PyErr_NoMemory();
3029 goto error;
3030 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003031
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003032 for (i=0 ; i<n ; i++)
3033 indices[i] = i;
3034 for (i=0 ; i<r ; i++)
3035 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003036
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003037 /* create permutationsobject structure */
3038 po = (permutationsobject *)type->tp_alloc(type, 0);
3039 if (po == NULL)
3040 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003041
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003042 po->pool = pool;
3043 po->indices = indices;
3044 po->cycles = cycles;
3045 po->result = NULL;
3046 po->r = r;
3047 po->stopped = r > n ? 1 : 0;
3048
3049 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003050
3051error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003052 if (indices != NULL)
3053 PyMem_Free(indices);
3054 if (cycles != NULL)
3055 PyMem_Free(cycles);
3056 Py_XDECREF(pool);
3057 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003058}
3059
3060static void
3061permutations_dealloc(permutationsobject *po)
3062{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003063 PyObject_GC_UnTrack(po);
3064 Py_XDECREF(po->pool);
3065 Py_XDECREF(po->result);
3066 PyMem_Free(po->indices);
3067 PyMem_Free(po->cycles);
3068 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003069}
3070
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003071static PyObject *
3072permutations_sizeof(permutationsobject *po, void *unused)
3073{
3074 Py_ssize_t res;
3075
3076 res = sizeof(permutationsobject);
3077 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3078 res += po->r * sizeof(Py_ssize_t);
3079 return PyLong_FromSsize_t(res);
3080}
3081
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003082static int
3083permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3084{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003085 Py_VISIT(po->pool);
3086 Py_VISIT(po->result);
3087 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003088}
3089
3090static PyObject *
3091permutations_next(permutationsobject *po)
3092{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003093 PyObject *elem;
3094 PyObject *oldelem;
3095 PyObject *pool = po->pool;
3096 Py_ssize_t *indices = po->indices;
3097 Py_ssize_t *cycles = po->cycles;
3098 PyObject *result = po->result;
3099 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3100 Py_ssize_t r = po->r;
3101 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003102
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003103 if (po->stopped)
3104 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003105
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003106 if (result == NULL) {
3107 /* On the first pass, initialize result tuple using the indices */
3108 result = PyTuple_New(r);
3109 if (result == NULL)
3110 goto empty;
3111 po->result = result;
3112 for (i=0; i<r ; i++) {
3113 index = indices[i];
3114 elem = PyTuple_GET_ITEM(pool, index);
3115 Py_INCREF(elem);
3116 PyTuple_SET_ITEM(result, i, elem);
3117 }
3118 } else {
3119 if (n == 0)
3120 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 /* Copy the previous result tuple or re-use it if available */
3123 if (Py_REFCNT(result) > 1) {
3124 PyObject *old_result = result;
3125 result = PyTuple_New(r);
3126 if (result == NULL)
3127 goto empty;
3128 po->result = result;
3129 for (i=0; i<r ; i++) {
3130 elem = PyTuple_GET_ITEM(old_result, i);
3131 Py_INCREF(elem);
3132 PyTuple_SET_ITEM(result, i, elem);
3133 }
3134 Py_DECREF(old_result);
3135 }
3136 /* Now, we've got the only copy so we can update it in-place */
3137 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003138
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003139 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3140 for (i=r-1 ; i>=0 ; i--) {
3141 cycles[i] -= 1;
3142 if (cycles[i] == 0) {
3143 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3144 index = indices[i];
3145 for (j=i ; j<n-1 ; j++)
3146 indices[j] = indices[j+1];
3147 indices[n-1] = index;
3148 cycles[i] = n - i;
3149 } else {
3150 j = cycles[i];
3151 index = indices[i];
3152 indices[i] = indices[n-j];
3153 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003154
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003155 for (k=i; k<r ; k++) {
3156 /* start with i, the leftmost element that changed */
3157 /* yield tuple(pool[k] for k in indices[:r]) */
3158 index = indices[k];
3159 elem = PyTuple_GET_ITEM(pool, index);
3160 Py_INCREF(elem);
3161 oldelem = PyTuple_GET_ITEM(result, k);
3162 PyTuple_SET_ITEM(result, k, elem);
3163 Py_DECREF(oldelem);
3164 }
3165 break;
3166 }
3167 }
3168 /* If i is negative, then the cycles have all
3169 rolled-over and we're done. */
3170 if (i < 0)
3171 goto empty;
3172 }
3173 Py_INCREF(result);
3174 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003175
3176empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003177 po->stopped = 1;
3178 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003179}
3180
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003181static PyObject *
3182permutations_reduce(permutationsobject *po)
3183{
3184 if (po->result == NULL) {
3185 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3186 } else if (po->stopped) {
3187 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3188 } else {
3189 PyObject *indices=NULL, *cycles=NULL;
3190 Py_ssize_t n, i;
3191
3192 /* we must pickle the indices and cycles and use them for setstate */
3193 n = PyTuple_GET_SIZE(po->pool);
3194 indices = PyTuple_New(n);
3195 if (indices == NULL)
3196 goto err;
3197 for (i=0; i<n; i++){
3198 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3199 if (!index)
3200 goto err;
3201 PyTuple_SET_ITEM(indices, i, index);
3202 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003203
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003204 cycles = PyTuple_New(po->r);
3205 if (cycles == NULL)
3206 goto err;
3207 for (i=0; i<po->r; i++)
3208 {
3209 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3210 if (!index)
3211 goto err;
3212 PyTuple_SET_ITEM(cycles, i, index);
3213 }
3214 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3215 po->pool, po->r,
3216 indices, cycles);
3217 err:
3218 Py_XDECREF(indices);
3219 Py_XDECREF(cycles);
3220 return NULL;
3221 }
3222}
3223
3224static PyObject *
3225permutations_setstate(permutationsobject *po, PyObject *state)
3226{
3227 PyObject *indices, *cycles, *result;
3228 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003229
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003230 if (!PyArg_ParseTuple(state, "O!O!",
3231 &PyTuple_Type, &indices,
3232 &PyTuple_Type, &cycles))
3233 return NULL;
3234
3235 n = PyTuple_GET_SIZE(po->pool);
3236 if (PyTuple_GET_SIZE(indices) != n ||
3237 PyTuple_GET_SIZE(cycles) != po->r)
3238 {
3239 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3240 return NULL;
3241 }
3242
3243 for (i=0; i<n; i++)
3244 {
3245 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3246 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3247 if (index < 0 && PyErr_Occurred())
3248 return NULL; /* not an integer */
3249 /* clamp the index */
3250 if (index < 0)
3251 index = 0;
3252 else if (index > n-1)
3253 index = n-1;
3254 po->indices[i] = index;
3255 }
3256
3257 for (i=0; i<po->r; i++)
3258 {
3259 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3260 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3261 if (index < 0 && PyErr_Occurred())
3262 return NULL; /* not an integer */
3263 if (index < 1)
3264 index = 1;
3265 else if (index > n-i)
3266 index = n-i;
3267 po->cycles[i] = index;
3268 }
3269 result = PyTuple_New(po->r);
3270 if (result == NULL)
3271 return NULL;
3272 for (i=0; i<po->r; i++) {
3273 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3274 Py_INCREF(element);
3275 PyTuple_SET_ITEM(result, i, element);
3276 }
3277 Py_CLEAR(po->result);
3278 po->result = result;
3279 Py_RETURN_NONE;
3280}
3281
3282static PyMethodDef permuations_methods[] = {
3283 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3284 reduce_doc},
3285 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3286 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003287 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3288 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003289 {NULL, NULL} /* sentinel */
3290};
3291
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003292PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003293"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003294\n\
3295Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003296permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003297
3298static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003299 PyVarObject_HEAD_INIT(NULL, 0)
3300 "itertools.permutations", /* tp_name */
3301 sizeof(permutationsobject), /* tp_basicsize */
3302 0, /* tp_itemsize */
3303 /* methods */
3304 (destructor)permutations_dealloc, /* tp_dealloc */
3305 0, /* tp_print */
3306 0, /* tp_getattr */
3307 0, /* tp_setattr */
3308 0, /* tp_reserved */
3309 0, /* tp_repr */
3310 0, /* tp_as_number */
3311 0, /* tp_as_sequence */
3312 0, /* tp_as_mapping */
3313 0, /* tp_hash */
3314 0, /* tp_call */
3315 0, /* tp_str */
3316 PyObject_GenericGetAttr, /* tp_getattro */
3317 0, /* tp_setattro */
3318 0, /* tp_as_buffer */
3319 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3320 Py_TPFLAGS_BASETYPE, /* tp_flags */
3321 permutations_doc, /* tp_doc */
3322 (traverseproc)permutations_traverse, /* tp_traverse */
3323 0, /* tp_clear */
3324 0, /* tp_richcompare */
3325 0, /* tp_weaklistoffset */
3326 PyObject_SelfIter, /* tp_iter */
3327 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003328 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003329 0, /* tp_members */
3330 0, /* tp_getset */
3331 0, /* tp_base */
3332 0, /* tp_dict */
3333 0, /* tp_descr_get */
3334 0, /* tp_descr_set */
3335 0, /* tp_dictoffset */
3336 0, /* tp_init */
3337 0, /* tp_alloc */
3338 permutations_new, /* tp_new */
3339 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003340};
3341
Raymond Hettinger482ba772010-12-01 22:48:00 +00003342/* accumulate object ************************************************************/
3343
3344typedef struct {
3345 PyObject_HEAD
3346 PyObject *total;
3347 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003348 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003349} accumulateobject;
3350
3351static PyTypeObject accumulate_type;
3352
3353static PyObject *
3354accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3355{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003356 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003357 PyObject *iterable;
3358 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003359 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003360 accumulateobject *lz;
3361
Raymond Hettinger5d446132011-03-27 18:52:10 -07003362 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3363 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003364 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003365
3366 /* Get iterator. */
3367 it = PyObject_GetIter(iterable);
3368 if (it == NULL)
3369 return NULL;
3370
Raymond Hettinger482ba772010-12-01 22:48:00 +00003371 /* create accumulateobject structure */
3372 lz = (accumulateobject *)type->tp_alloc(type, 0);
3373 if (lz == NULL) {
3374 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003375 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003376 }
3377
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003378 if (binop != Py_None) {
3379 Py_XINCREF(binop);
3380 lz->binop = binop;
3381 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003382 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003383 lz->it = it;
3384 return (PyObject *)lz;
3385}
3386
3387static void
3388accumulate_dealloc(accumulateobject *lz)
3389{
3390 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003391 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003392 Py_XDECREF(lz->total);
3393 Py_XDECREF(lz->it);
3394 Py_TYPE(lz)->tp_free(lz);
3395}
3396
3397static int
3398accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3399{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003400 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003401 Py_VISIT(lz->it);
3402 Py_VISIT(lz->total);
3403 return 0;
3404}
3405
3406static PyObject *
3407accumulate_next(accumulateobject *lz)
3408{
3409 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003410
Raymond Hettinger482ba772010-12-01 22:48:00 +00003411 val = PyIter_Next(lz->it);
3412 if (val == NULL)
3413 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003414
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003415 if (lz->total == NULL) {
3416 Py_INCREF(val);
3417 lz->total = val;
3418 return lz->total;
3419 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003420
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003421 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003422 newtotal = PyNumber_Add(lz->total, val);
3423 else
3424 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003425 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003426 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003427 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003428
3429 oldtotal = lz->total;
3430 lz->total = newtotal;
3431 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003432
Raymond Hettinger482ba772010-12-01 22:48:00 +00003433 Py_INCREF(newtotal);
3434 return newtotal;
3435}
3436
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003437static PyObject *
3438accumulate_reduce(accumulateobject *lz)
3439{
3440 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3441 lz->it, lz->binop?lz->binop:Py_None,
3442 lz->total?lz->total:Py_None);
3443 }
3444
3445static PyObject *
3446accumulate_setstate(accumulateobject *lz, PyObject *state)
3447{
3448 Py_CLEAR(lz->total);
3449 lz->total = state;
3450 Py_INCREF(lz->total);
3451 Py_RETURN_NONE;
3452}
3453
3454static PyMethodDef accumulate_methods[] = {
3455 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3456 reduce_doc},
3457 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3458 setstate_doc},
3459 {NULL, NULL} /* sentinel */
3460};
3461
Raymond Hettinger482ba772010-12-01 22:48:00 +00003462PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003463"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003464\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003465Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003466
3467static PyTypeObject accumulate_type = {
3468 PyVarObject_HEAD_INIT(NULL, 0)
3469 "itertools.accumulate", /* tp_name */
3470 sizeof(accumulateobject), /* tp_basicsize */
3471 0, /* tp_itemsize */
3472 /* methods */
3473 (destructor)accumulate_dealloc, /* tp_dealloc */
3474 0, /* tp_print */
3475 0, /* tp_getattr */
3476 0, /* tp_setattr */
3477 0, /* tp_reserved */
3478 0, /* tp_repr */
3479 0, /* tp_as_number */
3480 0, /* tp_as_sequence */
3481 0, /* tp_as_mapping */
3482 0, /* tp_hash */
3483 0, /* tp_call */
3484 0, /* tp_str */
3485 PyObject_GenericGetAttr, /* tp_getattro */
3486 0, /* tp_setattro */
3487 0, /* tp_as_buffer */
3488 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3489 Py_TPFLAGS_BASETYPE, /* tp_flags */
3490 accumulate_doc, /* tp_doc */
3491 (traverseproc)accumulate_traverse, /* tp_traverse */
3492 0, /* tp_clear */
3493 0, /* tp_richcompare */
3494 0, /* tp_weaklistoffset */
3495 PyObject_SelfIter, /* tp_iter */
3496 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003497 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003498 0, /* tp_members */
3499 0, /* tp_getset */
3500 0, /* tp_base */
3501 0, /* tp_dict */
3502 0, /* tp_descr_get */
3503 0, /* tp_descr_set */
3504 0, /* tp_dictoffset */
3505 0, /* tp_init */
3506 0, /* tp_alloc */
3507 accumulate_new, /* tp_new */
3508 PyObject_GC_Del, /* tp_free */
3509};
3510
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003511
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003512/* compress object ************************************************************/
3513
3514/* Equivalent to:
3515
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003516 def compress(data, selectors):
3517 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3518 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003519*/
3520
3521typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003522 PyObject_HEAD
3523 PyObject *data;
3524 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003525} compressobject;
3526
3527static PyTypeObject compress_type;
3528
3529static PyObject *
3530compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3531{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003532 PyObject *seq1, *seq2;
3533 PyObject *data=NULL, *selectors=NULL;
3534 compressobject *lz;
3535 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003536
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003537 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3538 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003539
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003540 data = PyObject_GetIter(seq1);
3541 if (data == NULL)
3542 goto fail;
3543 selectors = PyObject_GetIter(seq2);
3544 if (selectors == NULL)
3545 goto fail;
3546
3547 /* create compressobject structure */
3548 lz = (compressobject *)type->tp_alloc(type, 0);
3549 if (lz == NULL)
3550 goto fail;
3551 lz->data = data;
3552 lz->selectors = selectors;
3553 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003554
3555fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003556 Py_XDECREF(data);
3557 Py_XDECREF(selectors);
3558 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003559}
3560
3561static void
3562compress_dealloc(compressobject *lz)
3563{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003564 PyObject_GC_UnTrack(lz);
3565 Py_XDECREF(lz->data);
3566 Py_XDECREF(lz->selectors);
3567 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003568}
3569
3570static int
3571compress_traverse(compressobject *lz, visitproc visit, void *arg)
3572{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003573 Py_VISIT(lz->data);
3574 Py_VISIT(lz->selectors);
3575 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003576}
3577
3578static PyObject *
3579compress_next(compressobject *lz)
3580{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003581 PyObject *data = lz->data, *selectors = lz->selectors;
3582 PyObject *datum, *selector;
3583 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3584 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3585 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003586
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003587 while (1) {
3588 /* Steps: get datum, get selector, evaluate selector.
3589 Order is important (to match the pure python version
3590 in terms of which input gets a chance to raise an
3591 exception first).
3592 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003593
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003594 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003595 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003596 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003598 selector = selectornext(selectors);
3599 if (selector == NULL) {
3600 Py_DECREF(datum);
3601 return NULL;
3602 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003603
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003604 ok = PyObject_IsTrue(selector);
3605 Py_DECREF(selector);
3606 if (ok == 1)
3607 return datum;
3608 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003609 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003610 return NULL;
3611 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003612}
3613
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003614static PyObject *
3615compress_reduce(compressobject *lz)
3616{
3617 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3618 lz->data, lz->selectors);
3619 }
3620
3621static PyMethodDef compress_methods[] = {
3622 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3623 reduce_doc},
3624 {NULL, NULL} /* sentinel */
3625};
3626
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003627PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003628"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003629\n\
3630Return data elements corresponding to true selector elements.\n\
3631Forms a shorter iterator from selected data elements using the\n\
3632selectors to choose the data elements.");
3633
3634static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003635 PyVarObject_HEAD_INIT(NULL, 0)
3636 "itertools.compress", /* tp_name */
3637 sizeof(compressobject), /* tp_basicsize */
3638 0, /* tp_itemsize */
3639 /* methods */
3640 (destructor)compress_dealloc, /* tp_dealloc */
3641 0, /* tp_print */
3642 0, /* tp_getattr */
3643 0, /* tp_setattr */
3644 0, /* tp_reserved */
3645 0, /* tp_repr */
3646 0, /* tp_as_number */
3647 0, /* tp_as_sequence */
3648 0, /* tp_as_mapping */
3649 0, /* tp_hash */
3650 0, /* tp_call */
3651 0, /* tp_str */
3652 PyObject_GenericGetAttr, /* tp_getattro */
3653 0, /* tp_setattro */
3654 0, /* tp_as_buffer */
3655 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3656 Py_TPFLAGS_BASETYPE, /* tp_flags */
3657 compress_doc, /* tp_doc */
3658 (traverseproc)compress_traverse, /* tp_traverse */
3659 0, /* tp_clear */
3660 0, /* tp_richcompare */
3661 0, /* tp_weaklistoffset */
3662 PyObject_SelfIter, /* tp_iter */
3663 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003664 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003665 0, /* tp_members */
3666 0, /* tp_getset */
3667 0, /* tp_base */
3668 0, /* tp_dict */
3669 0, /* tp_descr_get */
3670 0, /* tp_descr_set */
3671 0, /* tp_dictoffset */
3672 0, /* tp_init */
3673 0, /* tp_alloc */
3674 compress_new, /* tp_new */
3675 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003676};
3677
3678
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003679/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003680
3681typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003682 PyObject_HEAD
3683 PyObject *func;
3684 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003685} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003686
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003687static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003688
3689static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003690filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003691{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003692 PyObject *func, *seq;
3693 PyObject *it;
3694 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003695
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003696 if (type == &filterfalse_type &&
3697 !_PyArg_NoKeywords("filterfalse()", kwds))
3698 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003699
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003700 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3701 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003702
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003703 /* Get iterator. */
3704 it = PyObject_GetIter(seq);
3705 if (it == NULL)
3706 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003707
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003708 /* create filterfalseobject structure */
3709 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3710 if (lz == NULL) {
3711 Py_DECREF(it);
3712 return NULL;
3713 }
3714 Py_INCREF(func);
3715 lz->func = func;
3716 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003718 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003719}
3720
3721static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003722filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003723{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003724 PyObject_GC_UnTrack(lz);
3725 Py_XDECREF(lz->func);
3726 Py_XDECREF(lz->it);
3727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003728}
3729
3730static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003731filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003732{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003733 Py_VISIT(lz->it);
3734 Py_VISIT(lz->func);
3735 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003736}
3737
3738static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003739filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003740{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003741 PyObject *item;
3742 PyObject *it = lz->it;
3743 long ok;
3744 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003746 iternext = *Py_TYPE(it)->tp_iternext;
3747 for (;;) {
3748 item = iternext(it);
3749 if (item == NULL)
3750 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003752 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3753 ok = PyObject_IsTrue(item);
3754 } else {
3755 PyObject *good;
3756 good = PyObject_CallFunctionObjArgs(lz->func,
3757 item, NULL);
3758 if (good == NULL) {
3759 Py_DECREF(item);
3760 return NULL;
3761 }
3762 ok = PyObject_IsTrue(good);
3763 Py_DECREF(good);
3764 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003765 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003766 return item;
3767 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003768 if (ok < 0)
3769 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003770 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003771}
3772
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003773static PyObject *
3774filterfalse_reduce(filterfalseobject *lz)
3775{
3776 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3777 lz->func, lz->it);
3778 }
3779
3780static PyMethodDef filterfalse_methods[] = {
3781 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3782 reduce_doc},
3783 {NULL, NULL} /* sentinel */
3784};
3785
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003786PyDoc_STRVAR(filterfalse_doc,
3787"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003788\n\
3789Return those items of sequence for which function(item) is false.\n\
3790If function is None, return the items that are false.");
3791
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003792static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003793 PyVarObject_HEAD_INIT(NULL, 0)
3794 "itertools.filterfalse", /* tp_name */
3795 sizeof(filterfalseobject), /* tp_basicsize */
3796 0, /* tp_itemsize */
3797 /* methods */
3798 (destructor)filterfalse_dealloc, /* tp_dealloc */
3799 0, /* tp_print */
3800 0, /* tp_getattr */
3801 0, /* tp_setattr */
3802 0, /* tp_reserved */
3803 0, /* tp_repr */
3804 0, /* tp_as_number */
3805 0, /* tp_as_sequence */
3806 0, /* tp_as_mapping */
3807 0, /* tp_hash */
3808 0, /* tp_call */
3809 0, /* tp_str */
3810 PyObject_GenericGetAttr, /* tp_getattro */
3811 0, /* tp_setattro */
3812 0, /* tp_as_buffer */
3813 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3814 Py_TPFLAGS_BASETYPE, /* tp_flags */
3815 filterfalse_doc, /* tp_doc */
3816 (traverseproc)filterfalse_traverse, /* tp_traverse */
3817 0, /* tp_clear */
3818 0, /* tp_richcompare */
3819 0, /* tp_weaklistoffset */
3820 PyObject_SelfIter, /* tp_iter */
3821 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003822 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003823 0, /* tp_members */
3824 0, /* tp_getset */
3825 0, /* tp_base */
3826 0, /* tp_dict */
3827 0, /* tp_descr_get */
3828 0, /* tp_descr_set */
3829 0, /* tp_dictoffset */
3830 0, /* tp_init */
3831 0, /* tp_alloc */
3832 filterfalse_new, /* tp_new */
3833 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003834};
3835
3836
3837/* count object ************************************************************/
3838
3839typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003840 PyObject_HEAD
3841 Py_ssize_t cnt;
3842 PyObject *long_cnt;
3843 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003844} countobject;
3845
Raymond Hettinger30729212009-02-12 06:28:27 +00003846/* Counting logic and invariants:
3847
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003848fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003850 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3851 Advances with: cnt += 1
3852 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003853
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003854slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003855
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003856 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3857 All counting is done with python objects (no overflows or underflows).
3858 Advances with: long_cnt += long_step
3859 Step may be zero -- effectively a slow version of repeat(cnt).
3860 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003861*/
3862
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003863static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003864
3865static PyObject *
3866count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3867{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003868 countobject *lz;
3869 int slow_mode = 0;
3870 Py_ssize_t cnt = 0;
3871 PyObject *long_cnt = NULL;
3872 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003873 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003874 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003875
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003876 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3877 kwlist, &long_cnt, &long_step))
3878 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003880 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3881 (long_step != NULL && !PyNumber_Check(long_step))) {
3882 PyErr_SetString(PyExc_TypeError, "a number is required");
3883 return NULL;
3884 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003885
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003886 if (long_cnt != NULL) {
3887 cnt = PyLong_AsSsize_t(long_cnt);
3888 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3889 PyErr_Clear();
3890 slow_mode = 1;
3891 }
3892 Py_INCREF(long_cnt);
3893 } else {
3894 cnt = 0;
3895 long_cnt = PyLong_FromLong(0);
3896 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003897
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003898 /* If not specified, step defaults to 1 */
3899 if (long_step == NULL) {
3900 long_step = PyLong_FromLong(1);
3901 if (long_step == NULL) {
3902 Py_DECREF(long_cnt);
3903 return NULL;
3904 }
3905 } else
3906 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003907
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003908 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003909
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003910 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003911 step = PyLong_AsLong(long_step);
3912 if (step != 1) {
3913 slow_mode = 1;
3914 if (step == -1 && PyErr_Occurred())
3915 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003916 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003917
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003918 if (slow_mode)
3919 cnt = PY_SSIZE_T_MAX;
3920 else
3921 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003922
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003923 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3924 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3925 assert(slow_mode ||
3926 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003927
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003928 /* create countobject structure */
3929 lz = (countobject *)type->tp_alloc(type, 0);
3930 if (lz == NULL) {
3931 Py_XDECREF(long_cnt);
3932 return NULL;
3933 }
3934 lz->cnt = cnt;
3935 lz->long_cnt = long_cnt;
3936 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003937
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003938 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003939}
3940
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003941static void
3942count_dealloc(countobject *lz)
3943{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003944 PyObject_GC_UnTrack(lz);
3945 Py_XDECREF(lz->long_cnt);
3946 Py_XDECREF(lz->long_step);
3947 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003948}
3949
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003950static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003951count_traverse(countobject *lz, visitproc visit, void *arg)
3952{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003953 Py_VISIT(lz->long_cnt);
3954 Py_VISIT(lz->long_step);
3955 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003956}
3957
3958static PyObject *
3959count_nextlong(countobject *lz)
3960{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003961 PyObject *long_cnt;
3962 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003963
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003964 long_cnt = lz->long_cnt;
3965 if (long_cnt == NULL) {
3966 /* Switch to slow_mode */
3967 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3968 if (long_cnt == NULL)
3969 return NULL;
3970 }
3971 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003972
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003973 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3974 if (stepped_up == NULL)
3975 return NULL;
3976 lz->long_cnt = stepped_up;
3977 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003978}
3979
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003980static PyObject *
3981count_next(countobject *lz)
3982{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003983 if (lz->cnt == PY_SSIZE_T_MAX)
3984 return count_nextlong(lz);
3985 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003986}
3987
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003988static PyObject *
3989count_repr(countobject *lz)
3990{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003991 if (lz->cnt != PY_SSIZE_T_MAX)
3992 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003993
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003994 if (PyLong_Check(lz->long_step)) {
3995 long step = PyLong_AsLong(lz->long_step);
3996 if (step == -1 && PyErr_Occurred()) {
3997 PyErr_Clear();
3998 }
3999 if (step == 1) {
4000 /* Don't display step when it is an integer equal to 1 */
4001 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
4002 }
4003 }
4004 return PyUnicode_FromFormat("count(%R, %R)",
4005 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004006}
4007
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004008static PyObject *
4009count_reduce(countobject *lz)
4010{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004011 if (lz->cnt == PY_SSIZE_T_MAX)
4012 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4013 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004014}
4015
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004016static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004017 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004018 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004019 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004020};
4021
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004022PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00004023 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004024\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004025Return a count object whose .__next__() method returns consecutive values.\n\
4026Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00004027 def count(firstval=0, step=1):\n\
Eli Benderskyc554f722013-09-03 06:37:19 -07004028 x = firstval\n\
4029 while 1:\n\
4030 yield x\n\
4031 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004032
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004033static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004034 PyVarObject_HEAD_INIT(NULL, 0)
4035 "itertools.count", /* tp_name */
4036 sizeof(countobject), /* tp_basicsize */
4037 0, /* tp_itemsize */
4038 /* methods */
4039 (destructor)count_dealloc, /* tp_dealloc */
4040 0, /* tp_print */
4041 0, /* tp_getattr */
4042 0, /* tp_setattr */
4043 0, /* tp_reserved */
4044 (reprfunc)count_repr, /* tp_repr */
4045 0, /* tp_as_number */
4046 0, /* tp_as_sequence */
4047 0, /* tp_as_mapping */
4048 0, /* tp_hash */
4049 0, /* tp_call */
4050 0, /* tp_str */
4051 PyObject_GenericGetAttr, /* tp_getattro */
4052 0, /* tp_setattro */
4053 0, /* tp_as_buffer */
4054 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4055 Py_TPFLAGS_BASETYPE, /* tp_flags */
4056 count_doc, /* tp_doc */
4057 (traverseproc)count_traverse, /* tp_traverse */
4058 0, /* tp_clear */
4059 0, /* tp_richcompare */
4060 0, /* tp_weaklistoffset */
4061 PyObject_SelfIter, /* tp_iter */
4062 (iternextfunc)count_next, /* tp_iternext */
4063 count_methods, /* tp_methods */
4064 0, /* tp_members */
4065 0, /* tp_getset */
4066 0, /* tp_base */
4067 0, /* tp_dict */
4068 0, /* tp_descr_get */
4069 0, /* tp_descr_set */
4070 0, /* tp_dictoffset */
4071 0, /* tp_init */
4072 0, /* tp_alloc */
4073 count_new, /* tp_new */
4074 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004075};
4076
4077
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004078/* repeat object ************************************************************/
4079
4080typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004081 PyObject_HEAD
4082 PyObject *element;
4083 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004084} repeatobject;
4085
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004086static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004087
4088static PyObject *
4089repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4090{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004091 repeatobject *ro;
4092 PyObject *element;
4093 Py_ssize_t cnt = -1;
4094 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004096 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4097 &element, &cnt))
4098 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004100 if (PyTuple_Size(args) == 2 && cnt < 0)
4101 cnt = 0;
4102
4103 ro = (repeatobject *)type->tp_alloc(type, 0);
4104 if (ro == NULL)
4105 return NULL;
4106 Py_INCREF(element);
4107 ro->element = element;
4108 ro->cnt = cnt;
4109 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004110}
4111
4112static void
4113repeat_dealloc(repeatobject *ro)
4114{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004115 PyObject_GC_UnTrack(ro);
4116 Py_XDECREF(ro->element);
4117 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004118}
4119
4120static int
4121repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4122{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 Py_VISIT(ro->element);
4124 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004125}
4126
4127static PyObject *
4128repeat_next(repeatobject *ro)
4129{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004130 if (ro->cnt == 0)
4131 return NULL;
4132 if (ro->cnt > 0)
4133 ro->cnt--;
4134 Py_INCREF(ro->element);
4135 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004136}
4137
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004138static PyObject *
4139repeat_repr(repeatobject *ro)
4140{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004141 if (ro->cnt == -1)
4142 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4143 else
4144 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4145}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004146
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004147static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004148repeat_len(repeatobject *ro)
4149{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004150 if (ro->cnt == -1) {
4151 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4152 return NULL;
4153 }
4154 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004155}
4156
Armin Rigof5b3e362006-02-11 21:32:43 +00004157PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004158
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004159static PyObject *
4160repeat_reduce(repeatobject *ro)
4161{
4162 /* unpickle this so that a new repeat iterator is constructed with an
4163 * object, then call __setstate__ on it to set cnt
4164 */
4165 if (ro->cnt >= 0)
4166 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4167 else
4168 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4169}
4170
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004171static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004172 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004173 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004174 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004175};
4176
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004177PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004178"repeat(object [,times]) -> create an iterator which returns the object\n\
4179for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004180endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004181
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004182static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004183 PyVarObject_HEAD_INIT(NULL, 0)
4184 "itertools.repeat", /* tp_name */
4185 sizeof(repeatobject), /* tp_basicsize */
4186 0, /* tp_itemsize */
4187 /* methods */
4188 (destructor)repeat_dealloc, /* tp_dealloc */
4189 0, /* tp_print */
4190 0, /* tp_getattr */
4191 0, /* tp_setattr */
4192 0, /* tp_reserved */
4193 (reprfunc)repeat_repr, /* tp_repr */
4194 0, /* tp_as_number */
4195 0, /* tp_as_sequence */
4196 0, /* tp_as_mapping */
4197 0, /* tp_hash */
4198 0, /* tp_call */
4199 0, /* tp_str */
4200 PyObject_GenericGetAttr, /* tp_getattro */
4201 0, /* tp_setattro */
4202 0, /* tp_as_buffer */
4203 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4204 Py_TPFLAGS_BASETYPE, /* tp_flags */
4205 repeat_doc, /* tp_doc */
4206 (traverseproc)repeat_traverse, /* tp_traverse */
4207 0, /* tp_clear */
4208 0, /* tp_richcompare */
4209 0, /* tp_weaklistoffset */
4210 PyObject_SelfIter, /* tp_iter */
4211 (iternextfunc)repeat_next, /* tp_iternext */
4212 repeat_methods, /* tp_methods */
4213 0, /* tp_members */
4214 0, /* tp_getset */
4215 0, /* tp_base */
4216 0, /* tp_dict */
4217 0, /* tp_descr_get */
4218 0, /* tp_descr_set */
4219 0, /* tp_dictoffset */
4220 0, /* tp_init */
4221 0, /* tp_alloc */
4222 repeat_new, /* tp_new */
4223 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004224};
4225
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004226/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004227
4228#include "Python.h"
4229
4230typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004231 PyObject_HEAD
4232 Py_ssize_t tuplesize;
4233 Py_ssize_t numactive;
4234 PyObject *ittuple; /* tuple of iterators */
4235 PyObject *result;
4236 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004237} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004238
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004239static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004240
4241static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004242zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004243{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004244 ziplongestobject *lz;
4245 Py_ssize_t i;
4246 PyObject *ittuple; /* tuple of iterators */
4247 PyObject *result;
4248 PyObject *fillvalue = Py_None;
4249 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004250
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004251 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4252 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4253 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4254 PyErr_SetString(PyExc_TypeError,
4255 "zip_longest() got an unexpected keyword argument");
4256 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004257 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004258 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004259
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004260 /* args must be a tuple */
4261 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004262
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004263 /* obtain iterators */
4264 ittuple = PyTuple_New(tuplesize);
4265 if (ittuple == NULL)
4266 return NULL;
4267 for (i=0; i < tuplesize; ++i) {
4268 PyObject *item = PyTuple_GET_ITEM(args, i);
4269 PyObject *it = PyObject_GetIter(item);
4270 if (it == NULL) {
4271 if (PyErr_ExceptionMatches(PyExc_TypeError))
4272 PyErr_Format(PyExc_TypeError,
4273 "zip_longest argument #%zd must support iteration",
4274 i+1);
4275 Py_DECREF(ittuple);
4276 return NULL;
4277 }
4278 PyTuple_SET_ITEM(ittuple, i, it);
4279 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004280
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004281 /* create a result holder */
4282 result = PyTuple_New(tuplesize);
4283 if (result == NULL) {
4284 Py_DECREF(ittuple);
4285 return NULL;
4286 }
4287 for (i=0 ; i < tuplesize ; i++) {
4288 Py_INCREF(Py_None);
4289 PyTuple_SET_ITEM(result, i, Py_None);
4290 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004291
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004292 /* create ziplongestobject structure */
4293 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4294 if (lz == NULL) {
4295 Py_DECREF(ittuple);
4296 Py_DECREF(result);
4297 return NULL;
4298 }
4299 lz->ittuple = ittuple;
4300 lz->tuplesize = tuplesize;
4301 lz->numactive = tuplesize;
4302 lz->result = result;
4303 Py_INCREF(fillvalue);
4304 lz->fillvalue = fillvalue;
4305 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004306}
4307
4308static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004309zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004310{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004311 PyObject_GC_UnTrack(lz);
4312 Py_XDECREF(lz->ittuple);
4313 Py_XDECREF(lz->result);
4314 Py_XDECREF(lz->fillvalue);
4315 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004316}
4317
4318static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004319zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004320{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004321 Py_VISIT(lz->ittuple);
4322 Py_VISIT(lz->result);
4323 Py_VISIT(lz->fillvalue);
4324 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004325}
4326
4327static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004328zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004329{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004330 Py_ssize_t i;
4331 Py_ssize_t tuplesize = lz->tuplesize;
4332 PyObject *result = lz->result;
4333 PyObject *it;
4334 PyObject *item;
4335 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004336
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004337 if (tuplesize == 0)
4338 return NULL;
4339 if (lz->numactive == 0)
4340 return NULL;
4341 if (Py_REFCNT(result) == 1) {
4342 Py_INCREF(result);
4343 for (i=0 ; i < tuplesize ; i++) {
4344 it = PyTuple_GET_ITEM(lz->ittuple, i);
4345 if (it == NULL) {
4346 Py_INCREF(lz->fillvalue);
4347 item = lz->fillvalue;
4348 } else {
4349 item = PyIter_Next(it);
4350 if (item == NULL) {
4351 lz->numactive -= 1;
4352 if (lz->numactive == 0 || PyErr_Occurred()) {
4353 lz->numactive = 0;
4354 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004355 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004356 } else {
4357 Py_INCREF(lz->fillvalue);
4358 item = lz->fillvalue;
4359 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4360 Py_DECREF(it);
4361 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004362 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004363 }
4364 olditem = PyTuple_GET_ITEM(result, i);
4365 PyTuple_SET_ITEM(result, i, item);
4366 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004367 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004368 } else {
4369 result = PyTuple_New(tuplesize);
4370 if (result == NULL)
4371 return NULL;
4372 for (i=0 ; i < tuplesize ; i++) {
4373 it = PyTuple_GET_ITEM(lz->ittuple, i);
4374 if (it == NULL) {
4375 Py_INCREF(lz->fillvalue);
4376 item = lz->fillvalue;
4377 } else {
4378 item = PyIter_Next(it);
4379 if (item == NULL) {
4380 lz->numactive -= 1;
4381 if (lz->numactive == 0 || PyErr_Occurred()) {
4382 lz->numactive = 0;
4383 Py_DECREF(result);
4384 return NULL;
4385 } else {
4386 Py_INCREF(lz->fillvalue);
4387 item = lz->fillvalue;
4388 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4389 Py_DECREF(it);
4390 }
4391 }
4392 }
4393 PyTuple_SET_ITEM(result, i, item);
4394 }
4395 }
4396 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004397}
4398
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004399static PyObject *
4400zip_longest_reduce(ziplongestobject *lz)
4401{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004402
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004403 /* Create a new tuple with empty sequences where appropriate to pickle.
4404 * Then use setstate to set the fillvalue
4405 */
4406 int i;
4407 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4408 if (args == NULL)
4409 return NULL;
4410 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4411 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4412 if (elem == NULL) {
4413 elem = PyTuple_New(0);
4414 if (elem == NULL) {
4415 Py_DECREF(args);
4416 return NULL;
4417 }
4418 } else
4419 Py_INCREF(elem);
4420 PyTuple_SET_ITEM(args, i, elem);
4421 }
4422 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4423}
4424
4425static PyObject *
4426zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4427{
4428 Py_CLEAR(lz->fillvalue);
4429 lz->fillvalue = state;
4430 Py_INCREF(lz->fillvalue);
4431 Py_RETURN_NONE;
4432}
4433
4434static PyMethodDef zip_longest_methods[] = {
4435 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4436 reduce_doc},
4437 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4438 setstate_doc},
4439 {NULL, NULL} /* sentinel */
4440};
4441
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004442PyDoc_STRVAR(zip_longest_doc,
4443"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004444\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004445Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004446the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004447method continues until the longest iterable in the argument sequence\n\
4448is exhausted and then it raises StopIteration. When the shorter iterables\n\
4449are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4450defaults to None or can be specified by a keyword argument.\n\
4451");
4452
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004453static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004454 PyVarObject_HEAD_INIT(NULL, 0)
4455 "itertools.zip_longest", /* tp_name */
4456 sizeof(ziplongestobject), /* tp_basicsize */
4457 0, /* tp_itemsize */
4458 /* methods */
4459 (destructor)zip_longest_dealloc, /* tp_dealloc */
4460 0, /* tp_print */
4461 0, /* tp_getattr */
4462 0, /* tp_setattr */
4463 0, /* tp_reserved */
4464 0, /* tp_repr */
4465 0, /* tp_as_number */
4466 0, /* tp_as_sequence */
4467 0, /* tp_as_mapping */
4468 0, /* tp_hash */
4469 0, /* tp_call */
4470 0, /* tp_str */
4471 PyObject_GenericGetAttr, /* tp_getattro */
4472 0, /* tp_setattro */
4473 0, /* tp_as_buffer */
4474 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4475 Py_TPFLAGS_BASETYPE, /* tp_flags */
4476 zip_longest_doc, /* tp_doc */
4477 (traverseproc)zip_longest_traverse, /* tp_traverse */
4478 0, /* tp_clear */
4479 0, /* tp_richcompare */
4480 0, /* tp_weaklistoffset */
4481 PyObject_SelfIter, /* tp_iter */
4482 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004483 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004484 0, /* tp_members */
4485 0, /* tp_getset */
4486 0, /* tp_base */
4487 0, /* tp_dict */
4488 0, /* tp_descr_get */
4489 0, /* tp_descr_set */
4490 0, /* tp_dictoffset */
4491 0, /* tp_init */
4492 0, /* tp_alloc */
4493 zip_longest_new, /* tp_new */
4494 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004495};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004496
4497/* module level code ********************************************************/
4498
4499PyDoc_STRVAR(module_doc,
4500"Functional tools for creating and using iterators.\n\
4501\n\
4502Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004503count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004504cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004505repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004506\n\
4507Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004508accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004509chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger8df58f72013-09-09 01:29:40 -05004510chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004511compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4512dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4513groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004514filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004515islice(seq, [start,] stop [, step]) --> elements from\n\
4516 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004517starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004518tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004519takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004520zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004521\n\
4522Combinatoric generators:\n\
4523product(p, q, ... [repeat=1]) --> cartesian product\n\
4524permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004525combinations(p, r)\n\
4526combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004527");
4528
4529
Raymond Hettingerad983e72003-11-12 14:32:26 +00004530static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004531 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4532 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004533};
4534
Martin v. Löwis1a214512008-06-11 05:26:20 +00004535
4536static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004537 PyModuleDef_HEAD_INIT,
4538 "itertools",
4539 module_doc,
4540 -1,
4541 module_methods,
4542 NULL,
4543 NULL,
4544 NULL,
4545 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004546};
4547
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004548PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004549PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004550{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004551 int i;
4552 PyObject *m;
4553 char *name;
4554 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004555 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004556 &combinations_type,
4557 &cwr_type,
4558 &cycle_type,
4559 &dropwhile_type,
4560 &takewhile_type,
4561 &islice_type,
4562 &starmap_type,
4563 &chain_type,
4564 &compress_type,
4565 &filterfalse_type,
4566 &count_type,
4567 &ziplongest_type,
4568 &permutations_type,
4569 &product_type,
4570 &repeat_type,
4571 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004572 &_grouper_type,
4573 &tee_type,
4574 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004575 NULL
4576 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004577
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004578 Py_TYPE(&teedataobject_type) = &PyType_Type;
4579 m = PyModule_Create(&itertoolsmodule);
4580 if (m == NULL)
4581 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004582
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004583 for (i=0 ; typelist[i] != NULL ; i++) {
4584 if (PyType_Ready(typelist[i]) < 0)
4585 return NULL;
4586 name = strchr(typelist[i]->tp_name, '.');
4587 assert (name != NULL);
4588 Py_INCREF(typelist[i]);
4589 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4590 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004591
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004592 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004593}