blob: 1494f8d6784cfbcc604bddd688da4b5b5b7af5a6 [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>
7 Copyright (c) 2003 Python Software Foundation.
8 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;
404 int numread;
405 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;
412 int index;
413 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,
1557"islice(iterable, [start,] stop [, step]) --> islice object\n\
1558\n\
1559Return an iterator whose next() method returns selected values from an\n\
1560iterable. If start is specified, will skip all preceding elements;\n\
1561otherwise, start defaults to zero. Step defaults to one. If\n\
1562specified as another value, step determines how many values are \n\
1563skipped between successive calls. Works like a slice() on a list\n\
1564but returns an iterator.");
1565
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001566static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001567 PyVarObject_HEAD_INIT(NULL, 0)
1568 "itertools.islice", /* tp_name */
1569 sizeof(isliceobject), /* tp_basicsize */
1570 0, /* tp_itemsize */
1571 /* methods */
1572 (destructor)islice_dealloc, /* tp_dealloc */
1573 0, /* tp_print */
1574 0, /* tp_getattr */
1575 0, /* tp_setattr */
1576 0, /* tp_reserved */
1577 0, /* tp_repr */
1578 0, /* tp_as_number */
1579 0, /* tp_as_sequence */
1580 0, /* tp_as_mapping */
1581 0, /* tp_hash */
1582 0, /* tp_call */
1583 0, /* tp_str */
1584 PyObject_GenericGetAttr, /* tp_getattro */
1585 0, /* tp_setattro */
1586 0, /* tp_as_buffer */
1587 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1588 Py_TPFLAGS_BASETYPE, /* tp_flags */
1589 islice_doc, /* tp_doc */
1590 (traverseproc)islice_traverse, /* tp_traverse */
1591 0, /* tp_clear */
1592 0, /* tp_richcompare */
1593 0, /* tp_weaklistoffset */
1594 PyObject_SelfIter, /* tp_iter */
1595 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001596 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001597 0, /* tp_members */
1598 0, /* tp_getset */
1599 0, /* tp_base */
1600 0, /* tp_dict */
1601 0, /* tp_descr_get */
1602 0, /* tp_descr_set */
1603 0, /* tp_dictoffset */
1604 0, /* tp_init */
1605 0, /* tp_alloc */
1606 islice_new, /* tp_new */
1607 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608};
1609
1610
1611/* starmap object ************************************************************/
1612
1613typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001614 PyObject_HEAD
1615 PyObject *func;
1616 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617} starmapobject;
1618
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001619static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001620
1621static PyObject *
1622starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1623{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001624 PyObject *func, *seq;
1625 PyObject *it;
1626 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001627
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001628 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1629 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001630
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001631 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1632 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001633
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001634 /* Get iterator. */
1635 it = PyObject_GetIter(seq);
1636 if (it == NULL)
1637 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001639 /* create starmapobject structure */
1640 lz = (starmapobject *)type->tp_alloc(type, 0);
1641 if (lz == NULL) {
1642 Py_DECREF(it);
1643 return NULL;
1644 }
1645 Py_INCREF(func);
1646 lz->func = func;
1647 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001649 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650}
1651
1652static void
1653starmap_dealloc(starmapobject *lz)
1654{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001655 PyObject_GC_UnTrack(lz);
1656 Py_XDECREF(lz->func);
1657 Py_XDECREF(lz->it);
1658 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001659}
1660
1661static int
1662starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1663{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001664 Py_VISIT(lz->it);
1665 Py_VISIT(lz->func);
1666 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001667}
1668
1669static PyObject *
1670starmap_next(starmapobject *lz)
1671{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001672 PyObject *args;
1673 PyObject *result;
1674 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001676 args = (*Py_TYPE(it)->tp_iternext)(it);
1677 if (args == NULL)
1678 return NULL;
1679 if (!PyTuple_CheckExact(args)) {
1680 PyObject *newargs = PySequence_Tuple(args);
1681 Py_DECREF(args);
1682 if (newargs == NULL)
1683 return NULL;
1684 args = newargs;
1685 }
1686 result = PyObject_Call(lz->func, args, NULL);
1687 Py_DECREF(args);
1688 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001689}
1690
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001691static PyObject *
1692starmap_reduce(starmapobject *lz)
1693{
1694 /* Just pickle the iterator */
1695 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1696}
1697
1698static PyMethodDef starmap_methods[] = {
1699 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1700 reduce_doc},
1701 {NULL, NULL} /* sentinel */
1702};
1703
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001704PyDoc_STRVAR(starmap_doc,
1705"starmap(function, sequence) --> starmap object\n\
1706\n\
1707Return an iterator whose values are returned from the function evaluated\n\
1708with a argument tuple taken from the given sequence.");
1709
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001710static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001711 PyVarObject_HEAD_INIT(NULL, 0)
1712 "itertools.starmap", /* tp_name */
1713 sizeof(starmapobject), /* tp_basicsize */
1714 0, /* tp_itemsize */
1715 /* methods */
1716 (destructor)starmap_dealloc, /* tp_dealloc */
1717 0, /* tp_print */
1718 0, /* tp_getattr */
1719 0, /* tp_setattr */
1720 0, /* tp_reserved */
1721 0, /* tp_repr */
1722 0, /* tp_as_number */
1723 0, /* tp_as_sequence */
1724 0, /* tp_as_mapping */
1725 0, /* tp_hash */
1726 0, /* tp_call */
1727 0, /* tp_str */
1728 PyObject_GenericGetAttr, /* tp_getattro */
1729 0, /* tp_setattro */
1730 0, /* tp_as_buffer */
1731 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1732 Py_TPFLAGS_BASETYPE, /* tp_flags */
1733 starmap_doc, /* tp_doc */
1734 (traverseproc)starmap_traverse, /* tp_traverse */
1735 0, /* tp_clear */
1736 0, /* tp_richcompare */
1737 0, /* tp_weaklistoffset */
1738 PyObject_SelfIter, /* tp_iter */
1739 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001740 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001741 0, /* tp_members */
1742 0, /* tp_getset */
1743 0, /* tp_base */
1744 0, /* tp_dict */
1745 0, /* tp_descr_get */
1746 0, /* tp_descr_set */
1747 0, /* tp_dictoffset */
1748 0, /* tp_init */
1749 0, /* tp_alloc */
1750 starmap_new, /* tp_new */
1751 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001752};
1753
1754
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001755/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001756
1757typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001758 PyObject_HEAD
1759 PyObject *source; /* Iterator over input iterables */
1760 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001761} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001762
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001763static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001765static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001766chain_new_internal(PyTypeObject *type, PyObject *source)
1767{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001768 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001769
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001770 lz = (chainobject *)type->tp_alloc(type, 0);
1771 if (lz == NULL) {
1772 Py_DECREF(source);
1773 return NULL;
1774 }
1775
1776 lz->source = source;
1777 lz->active = NULL;
1778 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001779}
1780
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001781static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001782chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001783{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001784 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001785
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001786 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1787 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001788
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001789 source = PyObject_GetIter(args);
1790 if (source == NULL)
1791 return NULL;
1792
1793 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001794}
1795
1796static PyObject *
1797chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1798{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001799 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001800
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001801 source = PyObject_GetIter(arg);
1802 if (source == NULL)
1803 return NULL;
1804
1805 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806}
1807
1808static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001809chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001810{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001811 PyObject_GC_UnTrack(lz);
1812 Py_XDECREF(lz->active);
1813 Py_XDECREF(lz->source);
1814 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001815}
1816
Raymond Hettinger2012f172003-02-07 05:32:58 +00001817static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001818chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001819{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001820 Py_VISIT(lz->source);
1821 Py_VISIT(lz->active);
1822 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001823}
1824
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001825static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001826chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001827{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001828 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001829
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001830 if (lz->source == NULL)
1831 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001832
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001833 if (lz->active == NULL) {
1834 PyObject *iterable = PyIter_Next(lz->source);
1835 if (iterable == NULL) {
1836 Py_CLEAR(lz->source);
1837 return NULL; /* no more input sources */
1838 }
1839 lz->active = PyObject_GetIter(iterable);
1840 Py_DECREF(iterable);
1841 if (lz->active == NULL) {
1842 Py_CLEAR(lz->source);
1843 return NULL; /* input not iterable */
1844 }
1845 }
1846 item = PyIter_Next(lz->active);
1847 if (item != NULL)
1848 return item;
1849 if (PyErr_Occurred()) {
1850 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1851 PyErr_Clear();
1852 else
1853 return NULL; /* input raised an exception */
1854 }
1855 Py_CLEAR(lz->active);
1856 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001857}
1858
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001859static PyObject *
1860chain_reduce(chainobject *lz)
1861{
1862 if (lz->source) {
1863 /* we can't pickle function objects (itertools.from_iterable) so
1864 * we must use setstate to replace the iterable. One day we
1865 * will fix pickling of functions
1866 */
1867 if (lz->active) {
1868 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1869 } else {
1870 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1871 }
1872 } else {
1873 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1874 }
1875 return NULL;
1876}
1877
1878static PyObject *
1879chain_setstate(chainobject *lz, PyObject *state)
1880{
1881 PyObject *source, *active=NULL;
1882 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1883 return NULL;
1884
1885 Py_CLEAR(lz->source);
1886 lz->source = source;
1887 Py_INCREF(lz->source);
1888 Py_CLEAR(lz->active);
1889 lz->active = active;
1890 Py_XINCREF(lz->active);
1891 Py_RETURN_NONE;
1892}
1893
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001894PyDoc_STRVAR(chain_doc,
1895"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001896\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001897Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001898first iterable until it is exhausted, then elements from the next\n\
1899iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001900
Christian Heimesf16baeb2008-02-29 14:57:44 +00001901PyDoc_STRVAR(chain_from_iterable_doc,
1902"chain.from_iterable(iterable) --> chain object\n\
1903\n\
1904Alternate chain() contructor taking a single iterable argument\n\
1905that evaluates lazily.");
1906
1907static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001908 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1909 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001910 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1911 reduce_doc},
1912 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1913 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001914 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001915};
1916
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001917static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001918 PyVarObject_HEAD_INIT(NULL, 0)
1919 "itertools.chain", /* tp_name */
1920 sizeof(chainobject), /* tp_basicsize */
1921 0, /* tp_itemsize */
1922 /* methods */
1923 (destructor)chain_dealloc, /* tp_dealloc */
1924 0, /* tp_print */
1925 0, /* tp_getattr */
1926 0, /* tp_setattr */
1927 0, /* tp_reserved */
1928 0, /* tp_repr */
1929 0, /* tp_as_number */
1930 0, /* tp_as_sequence */
1931 0, /* tp_as_mapping */
1932 0, /* tp_hash */
1933 0, /* tp_call */
1934 0, /* tp_str */
1935 PyObject_GenericGetAttr, /* tp_getattro */
1936 0, /* tp_setattro */
1937 0, /* tp_as_buffer */
1938 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1939 Py_TPFLAGS_BASETYPE, /* tp_flags */
1940 chain_doc, /* tp_doc */
1941 (traverseproc)chain_traverse, /* tp_traverse */
1942 0, /* tp_clear */
1943 0, /* tp_richcompare */
1944 0, /* tp_weaklistoffset */
1945 PyObject_SelfIter, /* tp_iter */
1946 (iternextfunc)chain_next, /* tp_iternext */
1947 chain_methods, /* tp_methods */
1948 0, /* tp_members */
1949 0, /* tp_getset */
1950 0, /* tp_base */
1951 0, /* tp_dict */
1952 0, /* tp_descr_get */
1953 0, /* tp_descr_set */
1954 0, /* tp_dictoffset */
1955 0, /* tp_init */
1956 0, /* tp_alloc */
1957 chain_new, /* tp_new */
1958 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001959};
1960
1961
Christian Heimesc3f30c42008-02-22 16:37:40 +00001962/* product object ************************************************************/
1963
1964typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001965 PyObject_HEAD
1966 PyObject *pools; /* tuple of pool tuples */
1967 Py_ssize_t *indices; /* one index per pool */
1968 PyObject *result; /* most recently returned result tuple */
1969 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001970} productobject;
1971
1972static PyTypeObject product_type;
1973
1974static PyObject *
1975product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1976{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001977 productobject *lz;
1978 Py_ssize_t nargs, npools, repeat=1;
1979 PyObject *pools = NULL;
1980 Py_ssize_t *indices = NULL;
1981 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001982
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001983 if (kwds != NULL) {
1984 char *kwlist[] = {"repeat", 0};
1985 PyObject *tmpargs = PyTuple_New(0);
1986 if (tmpargs == NULL)
1987 return NULL;
1988 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1989 Py_DECREF(tmpargs);
1990 return NULL;
1991 }
1992 Py_DECREF(tmpargs);
1993 if (repeat < 0) {
1994 PyErr_SetString(PyExc_ValueError,
1995 "repeat argument cannot be negative");
1996 return NULL;
1997 }
1998 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001999
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002000 assert(PyTuple_Check(args));
2001 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
2002 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002003
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002004 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
2005 if (indices == NULL) {
2006 PyErr_NoMemory();
2007 goto error;
2008 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002010 pools = PyTuple_New(npools);
2011 if (pools == NULL)
2012 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002013
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002014 for (i=0; i < nargs ; ++i) {
2015 PyObject *item = PyTuple_GET_ITEM(args, i);
2016 PyObject *pool = PySequence_Tuple(item);
2017 if (pool == NULL)
2018 goto error;
2019 PyTuple_SET_ITEM(pools, i, pool);
2020 indices[i] = 0;
2021 }
2022 for ( ; i < npools; ++i) {
2023 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2024 Py_INCREF(pool);
2025 PyTuple_SET_ITEM(pools, i, pool);
2026 indices[i] = 0;
2027 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002028
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002029 /* create productobject structure */
2030 lz = (productobject *)type->tp_alloc(type, 0);
2031 if (lz == NULL)
2032 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002034 lz->pools = pools;
2035 lz->indices = indices;
2036 lz->result = NULL;
2037 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002038
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002039 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002040
2041error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002042 if (indices != NULL)
2043 PyMem_Free(indices);
2044 Py_XDECREF(pools);
2045 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002046}
2047
2048static void
2049product_dealloc(productobject *lz)
2050{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002051 PyObject_GC_UnTrack(lz);
2052 Py_XDECREF(lz->pools);
2053 Py_XDECREF(lz->result);
2054 if (lz->indices != NULL)
2055 PyMem_Free(lz->indices);
2056 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002057}
2058
2059static int
2060product_traverse(productobject *lz, visitproc visit, void *arg)
2061{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002062 Py_VISIT(lz->pools);
2063 Py_VISIT(lz->result);
2064 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002065}
2066
2067static PyObject *
2068product_next(productobject *lz)
2069{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002070 PyObject *pool;
2071 PyObject *elem;
2072 PyObject *oldelem;
2073 PyObject *pools = lz->pools;
2074 PyObject *result = lz->result;
2075 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2076 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002077
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002078 if (lz->stopped)
2079 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002081 if (result == NULL) {
2082 /* On the first pass, return an initial tuple filled with the
2083 first element from each pool. */
2084 result = PyTuple_New(npools);
2085 if (result == NULL)
2086 goto empty;
2087 lz->result = result;
2088 for (i=0; i < npools; i++) {
2089 pool = PyTuple_GET_ITEM(pools, i);
2090 if (PyTuple_GET_SIZE(pool) == 0)
2091 goto empty;
2092 elem = PyTuple_GET_ITEM(pool, 0);
2093 Py_INCREF(elem);
2094 PyTuple_SET_ITEM(result, i, elem);
2095 }
2096 } else {
2097 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002098
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002099 /* Copy the previous result tuple or re-use it if available */
2100 if (Py_REFCNT(result) > 1) {
2101 PyObject *old_result = result;
2102 result = PyTuple_New(npools);
2103 if (result == NULL)
2104 goto empty;
2105 lz->result = result;
2106 for (i=0; i < npools; i++) {
2107 elem = PyTuple_GET_ITEM(old_result, i);
2108 Py_INCREF(elem);
2109 PyTuple_SET_ITEM(result, i, elem);
2110 }
2111 Py_DECREF(old_result);
2112 }
2113 /* Now, we've got the only copy so we can update it in-place */
2114 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002115
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002116 /* Update the pool indices right-to-left. Only advance to the
2117 next pool when the previous one rolls-over */
2118 for (i=npools-1 ; i >= 0 ; i--) {
2119 pool = PyTuple_GET_ITEM(pools, i);
2120 indices[i]++;
2121 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2122 /* Roll-over and advance to next pool */
2123 indices[i] = 0;
2124 elem = PyTuple_GET_ITEM(pool, 0);
2125 Py_INCREF(elem);
2126 oldelem = PyTuple_GET_ITEM(result, i);
2127 PyTuple_SET_ITEM(result, i, elem);
2128 Py_DECREF(oldelem);
2129 } else {
2130 /* No rollover. Just increment and stop here. */
2131 elem = PyTuple_GET_ITEM(pool, indices[i]);
2132 Py_INCREF(elem);
2133 oldelem = PyTuple_GET_ITEM(result, i);
2134 PyTuple_SET_ITEM(result, i, elem);
2135 Py_DECREF(oldelem);
2136 break;
2137 }
2138 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002139
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002140 /* If i is negative, then the indices have all rolled-over
2141 and we're done. */
2142 if (i < 0)
2143 goto empty;
2144 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002145
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002146 Py_INCREF(result);
2147 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002148
2149empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002150 lz->stopped = 1;
2151 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002152}
2153
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002154static PyObject *
2155product_reduce(productobject *lz)
2156{
2157 if (lz->stopped) {
2158 return Py_BuildValue("O(())", Py_TYPE(lz));
2159 } else if (lz->result == NULL) {
2160 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2161 } else {
2162 PyObject *indices;
2163 Py_ssize_t n, i;
2164
2165 /* we must pickle the indices use them for setstate, and
2166 * additionally indicate that the iterator has started
2167 */
2168 n = PyTuple_GET_SIZE(lz->pools);
2169 indices = PyTuple_New(n);
2170 if (indices == NULL)
2171 return NULL;
2172 for (i=0; i<n; i++){
2173 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2174 if (!index) {
2175 Py_DECREF(indices);
2176 return NULL;
2177 }
2178 PyTuple_SET_ITEM(indices, i, index);
2179 }
2180 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2181 }
2182}
2183
2184static PyObject *
2185product_setstate(productobject *lz, PyObject *state)
2186{
2187 PyObject *result;
2188 Py_ssize_t n, i;
2189
2190 n = PyTuple_GET_SIZE(lz->pools);
2191 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2192 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2193 return NULL;
2194 }
2195 for (i=0; i<n; i++)
2196 {
2197 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2198 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2199 if (index < 0 && PyErr_Occurred())
2200 return NULL; /* not an integer */
2201 /* clamp the index */
2202 if (index < 0)
2203 index = 0;
2204 else if (index > n-1)
2205 index = n-1;
2206 lz->indices[i] = index;
2207 }
2208
2209 result = PyTuple_New(n);
2210 if (!result)
2211 return NULL;
2212 for (i=0; i<n; i++) {
2213 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2214 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2215 Py_INCREF(element);
2216 PyTuple_SET_ITEM(result, i, element);
2217 }
2218 Py_CLEAR(lz->result);
2219 lz->result = result;
2220 Py_RETURN_NONE;
2221}
2222
2223static PyMethodDef product_methods[] = {
2224 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2225 reduce_doc},
2226 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2227 setstate_doc},
2228 {NULL, NULL} /* sentinel */
2229};
2230
Christian Heimesc3f30c42008-02-22 16:37:40 +00002231PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002232"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002233\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002234Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002235For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2236The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2237cycle in a manner similar to an odometer (with the rightmost element changing\n\
2238on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002239To compute the product of an iterable with itself, specify the number\n\
2240of repetitions with the optional repeat keyword argument. For example,\n\
2241product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002242product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2243product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2244
2245static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002246 PyVarObject_HEAD_INIT(NULL, 0)
2247 "itertools.product", /* tp_name */
2248 sizeof(productobject), /* tp_basicsize */
2249 0, /* tp_itemsize */
2250 /* methods */
2251 (destructor)product_dealloc, /* tp_dealloc */
2252 0, /* tp_print */
2253 0, /* tp_getattr */
2254 0, /* tp_setattr */
2255 0, /* tp_reserved */
2256 0, /* tp_repr */
2257 0, /* tp_as_number */
2258 0, /* tp_as_sequence */
2259 0, /* tp_as_mapping */
2260 0, /* tp_hash */
2261 0, /* tp_call */
2262 0, /* tp_str */
2263 PyObject_GenericGetAttr, /* tp_getattro */
2264 0, /* tp_setattro */
2265 0, /* tp_as_buffer */
2266 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2267 Py_TPFLAGS_BASETYPE, /* tp_flags */
2268 product_doc, /* tp_doc */
2269 (traverseproc)product_traverse, /* tp_traverse */
2270 0, /* tp_clear */
2271 0, /* tp_richcompare */
2272 0, /* tp_weaklistoffset */
2273 PyObject_SelfIter, /* tp_iter */
2274 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002275 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002276 0, /* tp_members */
2277 0, /* tp_getset */
2278 0, /* tp_base */
2279 0, /* tp_dict */
2280 0, /* tp_descr_get */
2281 0, /* tp_descr_set */
2282 0, /* tp_dictoffset */
2283 0, /* tp_init */
2284 0, /* tp_alloc */
2285 product_new, /* tp_new */
2286 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002287};
2288
2289
Christian Heimes380f7f22008-02-28 11:19:05 +00002290/* combinations object ************************************************************/
2291
2292typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002293 PyObject_HEAD
2294 PyObject *pool; /* input converted to a tuple */
2295 Py_ssize_t *indices; /* one index per result element */
2296 PyObject *result; /* most recently returned result tuple */
2297 Py_ssize_t r; /* size of result tuple */
2298 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002299} combinationsobject;
2300
2301static PyTypeObject combinations_type;
2302
2303static PyObject *
2304combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2305{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002306 combinationsobject *co;
2307 Py_ssize_t n;
2308 Py_ssize_t r;
2309 PyObject *pool = NULL;
2310 PyObject *iterable = NULL;
2311 Py_ssize_t *indices = NULL;
2312 Py_ssize_t i;
2313 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002315 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2316 &iterable, &r))
2317 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002318
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002319 pool = PySequence_Tuple(iterable);
2320 if (pool == NULL)
2321 goto error;
2322 n = PyTuple_GET_SIZE(pool);
2323 if (r < 0) {
2324 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2325 goto error;
2326 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002327
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002328 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2329 if (indices == NULL) {
2330 PyErr_NoMemory();
2331 goto error;
2332 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002333
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002334 for (i=0 ; i<r ; i++)
2335 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002336
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002337 /* create combinationsobject structure */
2338 co = (combinationsobject *)type->tp_alloc(type, 0);
2339 if (co == NULL)
2340 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002341
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002342 co->pool = pool;
2343 co->indices = indices;
2344 co->result = NULL;
2345 co->r = r;
2346 co->stopped = r > n ? 1 : 0;
2347
2348 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002349
2350error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002351 if (indices != NULL)
2352 PyMem_Free(indices);
2353 Py_XDECREF(pool);
2354 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002355}
2356
2357static void
2358combinations_dealloc(combinationsobject *co)
2359{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002360 PyObject_GC_UnTrack(co);
2361 Py_XDECREF(co->pool);
2362 Py_XDECREF(co->result);
2363 if (co->indices != NULL)
2364 PyMem_Free(co->indices);
2365 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002366}
2367
2368static int
2369combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2370{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002371 Py_VISIT(co->pool);
2372 Py_VISIT(co->result);
2373 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002374}
2375
2376static PyObject *
2377combinations_next(combinationsobject *co)
2378{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002379 PyObject *elem;
2380 PyObject *oldelem;
2381 PyObject *pool = co->pool;
2382 Py_ssize_t *indices = co->indices;
2383 PyObject *result = co->result;
2384 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2385 Py_ssize_t r = co->r;
2386 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002387
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002388 if (co->stopped)
2389 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002390
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002391 if (result == NULL) {
2392 /* On the first pass, initialize result tuple using the indices */
2393 result = PyTuple_New(r);
2394 if (result == NULL)
2395 goto empty;
2396 co->result = result;
2397 for (i=0; i<r ; i++) {
2398 index = indices[i];
2399 elem = PyTuple_GET_ITEM(pool, index);
2400 Py_INCREF(elem);
2401 PyTuple_SET_ITEM(result, i, elem);
2402 }
2403 } else {
2404 /* Copy the previous result tuple or re-use it if available */
2405 if (Py_REFCNT(result) > 1) {
2406 PyObject *old_result = result;
2407 result = PyTuple_New(r);
2408 if (result == NULL)
2409 goto empty;
2410 co->result = result;
2411 for (i=0; i<r ; i++) {
2412 elem = PyTuple_GET_ITEM(old_result, i);
2413 Py_INCREF(elem);
2414 PyTuple_SET_ITEM(result, i, elem);
2415 }
2416 Py_DECREF(old_result);
2417 }
2418 /* Now, we've got the only copy so we can update it in-place
2419 * CPython's empty tuple is a singleton and cached in
2420 * PyTuple's freelist.
2421 */
2422 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002423
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002424 /* Scan indices right-to-left until finding one that is not
2425 at its maximum (i + n - r). */
2426 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2427 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002428
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002429 /* If i is negative, then the indices are all at
2430 their maximum value and we're done. */
2431 if (i < 0)
2432 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002433
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002434 /* Increment the current index which we know is not at its
2435 maximum. Then move back to the right setting each index
2436 to its lowest possible value (one higher than the index
2437 to its left -- this maintains the sort order invariant). */
2438 indices[i]++;
2439 for (j=i+1 ; j<r ; j++)
2440 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002441
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002442 /* Update the result tuple for the new indices
2443 starting with i, the leftmost index that changed */
2444 for ( ; i<r ; i++) {
2445 index = indices[i];
2446 elem = PyTuple_GET_ITEM(pool, index);
2447 Py_INCREF(elem);
2448 oldelem = PyTuple_GET_ITEM(result, i);
2449 PyTuple_SET_ITEM(result, i, elem);
2450 Py_DECREF(oldelem);
2451 }
2452 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002453
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002454 Py_INCREF(result);
2455 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002456
2457empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002458 co->stopped = 1;
2459 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002460}
2461
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002462static PyObject *
2463combinations_reduce(combinationsobject *lz)
2464{
2465 if (lz->result == NULL) {
2466 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2467 } else if (lz->stopped) {
2468 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2469 } else {
2470 PyObject *indices;
2471 Py_ssize_t i;
2472
2473 /* we must pickle the indices and use them for setstate */
2474 indices = PyTuple_New(lz->r);
2475 if (!indices)
2476 return NULL;
2477 for (i=0; i<lz->r; i++)
2478 {
2479 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2480 if (!index) {
2481 Py_DECREF(indices);
2482 return NULL;
2483 }
2484 PyTuple_SET_ITEM(indices, i, index);
2485 }
2486
2487 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2488 }
2489}
2490
2491static PyObject *
2492combinations_setstate(combinationsobject *lz, PyObject *state)
2493{
2494 PyObject *result;
2495 Py_ssize_t i;
2496 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2497
2498 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2499 {
2500 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2501 return NULL;
2502 }
2503
2504 for (i=0; i<lz->r; i++)
2505 {
2506 Py_ssize_t max;
2507 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2508 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2509 if (index == -1 && PyErr_Occurred())
2510 return NULL; /* not an integer */
2511 max = i + n - lz->r;
2512 /* clamp the index (beware of negative max) */
2513 if (index > max)
2514 index = max;
2515 if (index < 0)
2516 index = 0;
2517 lz->indices[i] = index;
2518 }
2519
2520 result = PyTuple_New(lz->r);
2521 if (result == NULL)
2522 return NULL;
2523 for (i=0; i<lz->r; i++) {
2524 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2525 Py_INCREF(element);
2526 PyTuple_SET_ITEM(result, i, element);
2527 }
2528
2529 Py_CLEAR(lz->result);
2530 lz->result = result;
2531 Py_RETURN_NONE;
2532}
2533
2534static PyMethodDef combinations_methods[] = {
2535 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2536 reduce_doc},
2537 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2538 setstate_doc},
2539 {NULL, NULL} /* sentinel */
2540};
2541
Christian Heimes380f7f22008-02-28 11:19:05 +00002542PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002543"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002544\n\
2545Return successive r-length combinations of elements in the iterable.\n\n\
2546combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2547
2548static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002549 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002550 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002551 sizeof(combinationsobject), /* tp_basicsize */
2552 0, /* tp_itemsize */
2553 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002554 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002555 0, /* tp_print */
2556 0, /* tp_getattr */
2557 0, /* tp_setattr */
2558 0, /* tp_reserved */
2559 0, /* tp_repr */
2560 0, /* tp_as_number */
2561 0, /* tp_as_sequence */
2562 0, /* tp_as_mapping */
2563 0, /* tp_hash */
2564 0, /* tp_call */
2565 0, /* tp_str */
2566 PyObject_GenericGetAttr, /* tp_getattro */
2567 0, /* tp_setattro */
2568 0, /* tp_as_buffer */
2569 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2570 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002571 combinations_doc, /* tp_doc */
2572 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002573 0, /* tp_clear */
2574 0, /* tp_richcompare */
2575 0, /* tp_weaklistoffset */
2576 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002577 (iternextfunc)combinations_next, /* tp_iternext */
2578 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002579 0, /* tp_members */
2580 0, /* tp_getset */
2581 0, /* tp_base */
2582 0, /* tp_dict */
2583 0, /* tp_descr_get */
2584 0, /* tp_descr_set */
2585 0, /* tp_dictoffset */
2586 0, /* tp_init */
2587 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002588 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002589 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002590};
2591
2592
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002593/* combinations with replacement object *******************************************/
2594
2595/* Equivalent to:
2596
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002597 def combinations_with_replacement(iterable, r):
2598 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2599 # number items returned: (n+r-1)! / r! / (n-1)!
2600 pool = tuple(iterable)
2601 n = len(pool)
2602 indices = [0] * r
2603 yield tuple(pool[i] for i in indices)
2604 while 1:
2605 for i in reversed(range(r)):
2606 if indices[i] != n - 1:
2607 break
2608 else:
2609 return
2610 indices[i:] = [indices[i] + 1] * (r - i)
2611 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002612
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002613 def combinations_with_replacement2(iterable, r):
2614 'Alternate version that filters from product()'
2615 pool = tuple(iterable)
2616 n = len(pool)
2617 for indices in product(range(n), repeat=r):
2618 if sorted(indices) == list(indices):
2619 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002620*/
2621typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002622 PyObject_HEAD
2623 PyObject *pool; /* input converted to a tuple */
2624 Py_ssize_t *indices; /* one index per result element */
2625 PyObject *result; /* most recently returned result tuple */
2626 Py_ssize_t r; /* size of result tuple */
2627 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002628} cwrobject;
2629
2630static PyTypeObject cwr_type;
2631
2632static PyObject *
2633cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2634{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002635 cwrobject *co;
2636 Py_ssize_t n;
2637 Py_ssize_t r;
2638 PyObject *pool = NULL;
2639 PyObject *iterable = NULL;
2640 Py_ssize_t *indices = NULL;
2641 Py_ssize_t i;
2642 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002643
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002644 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2645 &iterable, &r))
2646 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002647
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002648 pool = PySequence_Tuple(iterable);
2649 if (pool == NULL)
2650 goto error;
2651 n = PyTuple_GET_SIZE(pool);
2652 if (r < 0) {
2653 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2654 goto error;
2655 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002656
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002657 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2658 if (indices == NULL) {
2659 PyErr_NoMemory();
2660 goto error;
2661 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002662
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002663 for (i=0 ; i<r ; i++)
2664 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002665
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002666 /* create cwrobject structure */
2667 co = (cwrobject *)type->tp_alloc(type, 0);
2668 if (co == NULL)
2669 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002670
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002671 co->pool = pool;
2672 co->indices = indices;
2673 co->result = NULL;
2674 co->r = r;
2675 co->stopped = !n && r;
2676
2677 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002678
2679error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002680 if (indices != NULL)
2681 PyMem_Free(indices);
2682 Py_XDECREF(pool);
2683 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002684}
2685
2686static void
2687cwr_dealloc(cwrobject *co)
2688{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002689 PyObject_GC_UnTrack(co);
2690 Py_XDECREF(co->pool);
2691 Py_XDECREF(co->result);
2692 if (co->indices != NULL)
2693 PyMem_Free(co->indices);
2694 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002695}
2696
2697static int
2698cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2699{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002700 Py_VISIT(co->pool);
2701 Py_VISIT(co->result);
2702 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002703}
2704
2705static PyObject *
2706cwr_next(cwrobject *co)
2707{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 PyObject *elem;
2709 PyObject *oldelem;
2710 PyObject *pool = co->pool;
2711 Py_ssize_t *indices = co->indices;
2712 PyObject *result = co->result;
2713 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2714 Py_ssize_t r = co->r;
2715 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002716
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002717 if (co->stopped)
2718 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002719
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002720 if (result == NULL) {
2721 /* On the first pass, initialize result tuple using the indices */
2722 result = PyTuple_New(r);
2723 if (result == NULL)
2724 goto empty;
2725 co->result = result;
2726 for (i=0; i<r ; i++) {
2727 index = indices[i];
2728 elem = PyTuple_GET_ITEM(pool, index);
2729 Py_INCREF(elem);
2730 PyTuple_SET_ITEM(result, i, elem);
2731 }
2732 } else {
2733 /* Copy the previous result tuple or re-use it if available */
2734 if (Py_REFCNT(result) > 1) {
2735 PyObject *old_result = result;
2736 result = PyTuple_New(r);
2737 if (result == NULL)
2738 goto empty;
2739 co->result = result;
2740 for (i=0; i<r ; i++) {
2741 elem = PyTuple_GET_ITEM(old_result, i);
2742 Py_INCREF(elem);
2743 PyTuple_SET_ITEM(result, i, elem);
2744 }
2745 Py_DECREF(old_result);
2746 }
2747 /* Now, we've got the only copy so we can update it in-place CPython's
2748 empty tuple is a singleton and cached in PyTuple's freelist. */
2749 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002751 /* Scan indices right-to-left until finding one that is not
2752 * at its maximum (n-1). */
2753 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2754 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 /* If i is negative, then the indices are all at
2757 their maximum value and we're done. */
2758 if (i < 0)
2759 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002760
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002761 /* Increment the current index which we know is not at its
2762 maximum. Then set all to the right to the same value. */
2763 indices[i]++;
2764 for (j=i+1 ; j<r ; j++)
2765 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002766
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002767 /* Update the result tuple for the new indices
2768 starting with i, the leftmost index that changed */
2769 for ( ; i<r ; i++) {
2770 index = indices[i];
2771 elem = PyTuple_GET_ITEM(pool, index);
2772 Py_INCREF(elem);
2773 oldelem = PyTuple_GET_ITEM(result, i);
2774 PyTuple_SET_ITEM(result, i, elem);
2775 Py_DECREF(oldelem);
2776 }
2777 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002778
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002779 Py_INCREF(result);
2780 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002781
2782empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002783 co->stopped = 1;
2784 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002785}
2786
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002787static PyObject *
2788cwr_reduce(cwrobject *lz)
2789{
2790 if (lz->result == NULL) {
2791 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2792 } else if (lz->stopped) {
2793 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2794 } else {
2795 PyObject *indices;
2796 Py_ssize_t i;
2797
2798 /* we must pickle the indices and use them for setstate */
2799 indices = PyTuple_New(lz->r);
2800 if (!indices)
2801 return NULL;
2802 for (i=0; i<lz->r; i++)
2803 {
2804 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2805 if (!index) {
2806 Py_DECREF(indices);
2807 return NULL;
2808 }
2809 PyTuple_SET_ITEM(indices, i, index);
2810 }
2811
2812 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2813 }
2814}
2815
2816static PyObject *
2817cwr_setstate(cwrobject *lz, PyObject *state)
2818{
2819 PyObject *result;
2820 Py_ssize_t n, i;
2821
2822 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2823 {
2824 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2825 return NULL;
2826 }
2827
2828 n = PyTuple_GET_SIZE(lz->pool);
2829 for (i=0; i<lz->r; i++)
2830 {
2831 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2832 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2833 if (index < 0 && PyErr_Occurred())
2834 return NULL; /* not an integer */
2835 /* clamp the index */
2836 if (index < 0)
2837 index = 0;
2838 else if (index > n-1)
2839 index = n-1;
2840 lz->indices[i] = index;
2841 }
2842 result = PyTuple_New(lz->r);
2843 if (result == NULL)
2844 return NULL;
2845 for (i=0; i<lz->r; i++) {
2846 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2847 Py_INCREF(element);
2848 PyTuple_SET_ITEM(result, i, element);
2849 }
2850 Py_CLEAR(lz->result);
2851 lz->result = result;
2852 Py_RETURN_NONE;
2853}
2854
2855static PyMethodDef cwr_methods[] = {
2856 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2857 reduce_doc},
2858 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2859 setstate_doc},
2860 {NULL, NULL} /* sentinel */
2861};
2862
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002863PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002864"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002865\n\
2866Return successive r-length combinations of elements in the iterable\n\
2867allowing individual elements to have successive repeats.\n\
2868combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2869
2870static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002871 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002872 "itertools.combinations_with_replacement", /* tp_name */
2873 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002874 0, /* tp_itemsize */
2875 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002876 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002877 0, /* tp_print */
2878 0, /* tp_getattr */
2879 0, /* tp_setattr */
2880 0, /* tp_reserved */
2881 0, /* tp_repr */
2882 0, /* tp_as_number */
2883 0, /* tp_as_sequence */
2884 0, /* tp_as_mapping */
2885 0, /* tp_hash */
2886 0, /* tp_call */
2887 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002888 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002889 0, /* tp_setattro */
2890 0, /* tp_as_buffer */
2891 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002892 Py_TPFLAGS_BASETYPE, /* tp_flags */
2893 cwr_doc, /* tp_doc */
2894 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002895 0, /* tp_clear */
2896 0, /* tp_richcompare */
2897 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002898 PyObject_SelfIter, /* tp_iter */
2899 (iternextfunc)cwr_next, /* tp_iternext */
2900 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002901 0, /* tp_members */
2902 0, /* tp_getset */
2903 0, /* tp_base */
2904 0, /* tp_dict */
2905 0, /* tp_descr_get */
2906 0, /* tp_descr_set */
2907 0, /* tp_dictoffset */
2908 0, /* tp_init */
2909 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002910 cwr_new, /* tp_new */
2911 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002912};
2913
2914
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002915/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002916
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002917def permutations(iterable, r=None):
2918 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2919 pool = tuple(iterable)
2920 n = len(pool)
2921 r = n if r is None else r
2922 indices = range(n)
2923 cycles = range(n-r+1, n+1)[::-1]
2924 yield tuple(pool[i] for i in indices[:r])
2925 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002926 for i in reversed(range(r)):
2927 cycles[i] -= 1
2928 if cycles[i] == 0:
2929 indices[i:] = indices[i+1:] + indices[i:i+1]
2930 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002931 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002932 j = cycles[i]
2933 indices[i], indices[-j] = indices[-j], indices[i]
2934 yield tuple(pool[i] for i in indices[:r])
2935 break
2936 else:
2937 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002938*/
2939
2940typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002941 PyObject_HEAD
2942 PyObject *pool; /* input converted to a tuple */
2943 Py_ssize_t *indices; /* one index per element in the pool */
2944 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2945 PyObject *result; /* most recently returned result tuple */
2946 Py_ssize_t r; /* size of result tuple */
2947 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002948} permutationsobject;
2949
2950static PyTypeObject permutations_type;
2951
2952static PyObject *
2953permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2954{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002955 permutationsobject *po;
2956 Py_ssize_t n;
2957 Py_ssize_t r;
2958 PyObject *robj = Py_None;
2959 PyObject *pool = NULL;
2960 PyObject *iterable = NULL;
2961 Py_ssize_t *indices = NULL;
2962 Py_ssize_t *cycles = NULL;
2963 Py_ssize_t i;
2964 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002965
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002966 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2967 &iterable, &robj))
2968 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002969
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002970 pool = PySequence_Tuple(iterable);
2971 if (pool == NULL)
2972 goto error;
2973 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002974
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002975 r = n;
2976 if (robj != Py_None) {
2977 if (!PyLong_Check(robj)) {
2978 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2979 goto error;
2980 }
2981 r = PyLong_AsSsize_t(robj);
2982 if (r == -1 && PyErr_Occurred())
2983 goto error;
2984 }
2985 if (r < 0) {
2986 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2987 goto error;
2988 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002989
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002990 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2991 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2992 if (indices == NULL || cycles == NULL) {
2993 PyErr_NoMemory();
2994 goto error;
2995 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002996
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002997 for (i=0 ; i<n ; i++)
2998 indices[i] = i;
2999 for (i=0 ; i<r ; i++)
3000 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003001
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003002 /* create permutationsobject structure */
3003 po = (permutationsobject *)type->tp_alloc(type, 0);
3004 if (po == NULL)
3005 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003007 po->pool = pool;
3008 po->indices = indices;
3009 po->cycles = cycles;
3010 po->result = NULL;
3011 po->r = r;
3012 po->stopped = r > n ? 1 : 0;
3013
3014 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003015
3016error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003017 if (indices != NULL)
3018 PyMem_Free(indices);
3019 if (cycles != NULL)
3020 PyMem_Free(cycles);
3021 Py_XDECREF(pool);
3022 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003023}
3024
3025static void
3026permutations_dealloc(permutationsobject *po)
3027{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003028 PyObject_GC_UnTrack(po);
3029 Py_XDECREF(po->pool);
3030 Py_XDECREF(po->result);
3031 PyMem_Free(po->indices);
3032 PyMem_Free(po->cycles);
3033 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003034}
3035
3036static int
3037permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3038{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003039 Py_VISIT(po->pool);
3040 Py_VISIT(po->result);
3041 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003042}
3043
3044static PyObject *
3045permutations_next(permutationsobject *po)
3046{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003047 PyObject *elem;
3048 PyObject *oldelem;
3049 PyObject *pool = po->pool;
3050 Py_ssize_t *indices = po->indices;
3051 Py_ssize_t *cycles = po->cycles;
3052 PyObject *result = po->result;
3053 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3054 Py_ssize_t r = po->r;
3055 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003057 if (po->stopped)
3058 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003059
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003060 if (result == NULL) {
3061 /* On the first pass, initialize result tuple using the indices */
3062 result = PyTuple_New(r);
3063 if (result == NULL)
3064 goto empty;
3065 po->result = result;
3066 for (i=0; i<r ; i++) {
3067 index = indices[i];
3068 elem = PyTuple_GET_ITEM(pool, index);
3069 Py_INCREF(elem);
3070 PyTuple_SET_ITEM(result, i, elem);
3071 }
3072 } else {
3073 if (n == 0)
3074 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003076 /* Copy the previous result tuple or re-use it if available */
3077 if (Py_REFCNT(result) > 1) {
3078 PyObject *old_result = result;
3079 result = PyTuple_New(r);
3080 if (result == NULL)
3081 goto empty;
3082 po->result = result;
3083 for (i=0; i<r ; i++) {
3084 elem = PyTuple_GET_ITEM(old_result, i);
3085 Py_INCREF(elem);
3086 PyTuple_SET_ITEM(result, i, elem);
3087 }
3088 Py_DECREF(old_result);
3089 }
3090 /* Now, we've got the only copy so we can update it in-place */
3091 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003093 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3094 for (i=r-1 ; i>=0 ; i--) {
3095 cycles[i] -= 1;
3096 if (cycles[i] == 0) {
3097 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3098 index = indices[i];
3099 for (j=i ; j<n-1 ; j++)
3100 indices[j] = indices[j+1];
3101 indices[n-1] = index;
3102 cycles[i] = n - i;
3103 } else {
3104 j = cycles[i];
3105 index = indices[i];
3106 indices[i] = indices[n-j];
3107 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003108
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003109 for (k=i; k<r ; k++) {
3110 /* start with i, the leftmost element that changed */
3111 /* yield tuple(pool[k] for k in indices[:r]) */
3112 index = indices[k];
3113 elem = PyTuple_GET_ITEM(pool, index);
3114 Py_INCREF(elem);
3115 oldelem = PyTuple_GET_ITEM(result, k);
3116 PyTuple_SET_ITEM(result, k, elem);
3117 Py_DECREF(oldelem);
3118 }
3119 break;
3120 }
3121 }
3122 /* If i is negative, then the cycles have all
3123 rolled-over and we're done. */
3124 if (i < 0)
3125 goto empty;
3126 }
3127 Py_INCREF(result);
3128 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003129
3130empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003131 po->stopped = 1;
3132 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003133}
3134
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003135static PyObject *
3136permutations_reduce(permutationsobject *po)
3137{
3138 if (po->result == NULL) {
3139 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3140 } else if (po->stopped) {
3141 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3142 } else {
3143 PyObject *indices=NULL, *cycles=NULL;
3144 Py_ssize_t n, i;
3145
3146 /* we must pickle the indices and cycles and use them for setstate */
3147 n = PyTuple_GET_SIZE(po->pool);
3148 indices = PyTuple_New(n);
3149 if (indices == NULL)
3150 goto err;
3151 for (i=0; i<n; i++){
3152 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3153 if (!index)
3154 goto err;
3155 PyTuple_SET_ITEM(indices, i, index);
3156 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003157
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003158 cycles = PyTuple_New(po->r);
3159 if (cycles == NULL)
3160 goto err;
3161 for (i=0; i<po->r; i++)
3162 {
3163 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3164 if (!index)
3165 goto err;
3166 PyTuple_SET_ITEM(cycles, i, index);
3167 }
3168 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3169 po->pool, po->r,
3170 indices, cycles);
3171 err:
3172 Py_XDECREF(indices);
3173 Py_XDECREF(cycles);
3174 return NULL;
3175 }
3176}
3177
3178static PyObject *
3179permutations_setstate(permutationsobject *po, PyObject *state)
3180{
3181 PyObject *indices, *cycles, *result;
3182 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003183
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003184 if (!PyArg_ParseTuple(state, "O!O!",
3185 &PyTuple_Type, &indices,
3186 &PyTuple_Type, &cycles))
3187 return NULL;
3188
3189 n = PyTuple_GET_SIZE(po->pool);
3190 if (PyTuple_GET_SIZE(indices) != n ||
3191 PyTuple_GET_SIZE(cycles) != po->r)
3192 {
3193 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3194 return NULL;
3195 }
3196
3197 for (i=0; i<n; i++)
3198 {
3199 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3200 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3201 if (index < 0 && PyErr_Occurred())
3202 return NULL; /* not an integer */
3203 /* clamp the index */
3204 if (index < 0)
3205 index = 0;
3206 else if (index > n-1)
3207 index = n-1;
3208 po->indices[i] = index;
3209 }
3210
3211 for (i=0; i<po->r; i++)
3212 {
3213 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3214 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3215 if (index < 0 && PyErr_Occurred())
3216 return NULL; /* not an integer */
3217 if (index < 1)
3218 index = 1;
3219 else if (index > n-i)
3220 index = n-i;
3221 po->cycles[i] = index;
3222 }
3223 result = PyTuple_New(po->r);
3224 if (result == NULL)
3225 return NULL;
3226 for (i=0; i<po->r; i++) {
3227 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3228 Py_INCREF(element);
3229 PyTuple_SET_ITEM(result, i, element);
3230 }
3231 Py_CLEAR(po->result);
3232 po->result = result;
3233 Py_RETURN_NONE;
3234}
3235
3236static PyMethodDef permuations_methods[] = {
3237 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3238 reduce_doc},
3239 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3240 setstate_doc},
3241 {NULL, NULL} /* sentinel */
3242};
3243
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003244PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003245"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003246\n\
3247Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003248permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003249
3250static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003251 PyVarObject_HEAD_INIT(NULL, 0)
3252 "itertools.permutations", /* tp_name */
3253 sizeof(permutationsobject), /* tp_basicsize */
3254 0, /* tp_itemsize */
3255 /* methods */
3256 (destructor)permutations_dealloc, /* tp_dealloc */
3257 0, /* tp_print */
3258 0, /* tp_getattr */
3259 0, /* tp_setattr */
3260 0, /* tp_reserved */
3261 0, /* tp_repr */
3262 0, /* tp_as_number */
3263 0, /* tp_as_sequence */
3264 0, /* tp_as_mapping */
3265 0, /* tp_hash */
3266 0, /* tp_call */
3267 0, /* tp_str */
3268 PyObject_GenericGetAttr, /* tp_getattro */
3269 0, /* tp_setattro */
3270 0, /* tp_as_buffer */
3271 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3272 Py_TPFLAGS_BASETYPE, /* tp_flags */
3273 permutations_doc, /* tp_doc */
3274 (traverseproc)permutations_traverse, /* tp_traverse */
3275 0, /* tp_clear */
3276 0, /* tp_richcompare */
3277 0, /* tp_weaklistoffset */
3278 PyObject_SelfIter, /* tp_iter */
3279 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003280 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003281 0, /* tp_members */
3282 0, /* tp_getset */
3283 0, /* tp_base */
3284 0, /* tp_dict */
3285 0, /* tp_descr_get */
3286 0, /* tp_descr_set */
3287 0, /* tp_dictoffset */
3288 0, /* tp_init */
3289 0, /* tp_alloc */
3290 permutations_new, /* tp_new */
3291 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003292};
3293
Raymond Hettinger482ba772010-12-01 22:48:00 +00003294/* accumulate object ************************************************************/
3295
3296typedef struct {
3297 PyObject_HEAD
3298 PyObject *total;
3299 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003300 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003301} accumulateobject;
3302
3303static PyTypeObject accumulate_type;
3304
3305static PyObject *
3306accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3307{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003308 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003309 PyObject *iterable;
3310 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003311 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003312 accumulateobject *lz;
3313
Raymond Hettinger5d446132011-03-27 18:52:10 -07003314 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3315 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003316 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003317
3318 /* Get iterator. */
3319 it = PyObject_GetIter(iterable);
3320 if (it == NULL)
3321 return NULL;
3322
Raymond Hettinger482ba772010-12-01 22:48:00 +00003323 /* create accumulateobject structure */
3324 lz = (accumulateobject *)type->tp_alloc(type, 0);
3325 if (lz == NULL) {
3326 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003327 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003328 }
3329
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003330 if (binop != Py_None) {
3331 Py_XINCREF(binop);
3332 lz->binop = binop;
3333 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003334 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003335 lz->it = it;
3336 return (PyObject *)lz;
3337}
3338
3339static void
3340accumulate_dealloc(accumulateobject *lz)
3341{
3342 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003343 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003344 Py_XDECREF(lz->total);
3345 Py_XDECREF(lz->it);
3346 Py_TYPE(lz)->tp_free(lz);
3347}
3348
3349static int
3350accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3351{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003352 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003353 Py_VISIT(lz->it);
3354 Py_VISIT(lz->total);
3355 return 0;
3356}
3357
3358static PyObject *
3359accumulate_next(accumulateobject *lz)
3360{
3361 PyObject *val, *oldtotal, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003362
Raymond Hettinger482ba772010-12-01 22:48:00 +00003363 val = PyIter_Next(lz->it);
3364 if (val == NULL)
3365 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003366
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003367 if (lz->total == NULL) {
3368 Py_INCREF(val);
3369 lz->total = val;
3370 return lz->total;
3371 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003372
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003373 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003374 newtotal = PyNumber_Add(lz->total, val);
3375 else
3376 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003377 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003378 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003379 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003380
3381 oldtotal = lz->total;
3382 lz->total = newtotal;
3383 Py_DECREF(oldtotal);
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003384
Raymond Hettinger482ba772010-12-01 22:48:00 +00003385 Py_INCREF(newtotal);
3386 return newtotal;
3387}
3388
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003389static PyObject *
3390accumulate_reduce(accumulateobject *lz)
3391{
3392 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3393 lz->it, lz->binop?lz->binop:Py_None,
3394 lz->total?lz->total:Py_None);
3395 }
3396
3397static PyObject *
3398accumulate_setstate(accumulateobject *lz, PyObject *state)
3399{
3400 Py_CLEAR(lz->total);
3401 lz->total = state;
3402 Py_INCREF(lz->total);
3403 Py_RETURN_NONE;
3404}
3405
3406static PyMethodDef accumulate_methods[] = {
3407 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3408 reduce_doc},
3409 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3410 setstate_doc},
3411 {NULL, NULL} /* sentinel */
3412};
3413
Raymond Hettinger482ba772010-12-01 22:48:00 +00003414PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003415"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003416\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003417Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003418
3419static PyTypeObject accumulate_type = {
3420 PyVarObject_HEAD_INIT(NULL, 0)
3421 "itertools.accumulate", /* tp_name */
3422 sizeof(accumulateobject), /* tp_basicsize */
3423 0, /* tp_itemsize */
3424 /* methods */
3425 (destructor)accumulate_dealloc, /* tp_dealloc */
3426 0, /* tp_print */
3427 0, /* tp_getattr */
3428 0, /* tp_setattr */
3429 0, /* tp_reserved */
3430 0, /* tp_repr */
3431 0, /* tp_as_number */
3432 0, /* tp_as_sequence */
3433 0, /* tp_as_mapping */
3434 0, /* tp_hash */
3435 0, /* tp_call */
3436 0, /* tp_str */
3437 PyObject_GenericGetAttr, /* tp_getattro */
3438 0, /* tp_setattro */
3439 0, /* tp_as_buffer */
3440 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3441 Py_TPFLAGS_BASETYPE, /* tp_flags */
3442 accumulate_doc, /* tp_doc */
3443 (traverseproc)accumulate_traverse, /* tp_traverse */
3444 0, /* tp_clear */
3445 0, /* tp_richcompare */
3446 0, /* tp_weaklistoffset */
3447 PyObject_SelfIter, /* tp_iter */
3448 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003449 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003450 0, /* tp_members */
3451 0, /* tp_getset */
3452 0, /* tp_base */
3453 0, /* tp_dict */
3454 0, /* tp_descr_get */
3455 0, /* tp_descr_set */
3456 0, /* tp_dictoffset */
3457 0, /* tp_init */
3458 0, /* tp_alloc */
3459 accumulate_new, /* tp_new */
3460 PyObject_GC_Del, /* tp_free */
3461};
3462
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003463
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003464/* compress object ************************************************************/
3465
3466/* Equivalent to:
3467
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003468 def compress(data, selectors):
3469 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3470 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003471*/
3472
3473typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003474 PyObject_HEAD
3475 PyObject *data;
3476 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003477} compressobject;
3478
3479static PyTypeObject compress_type;
3480
3481static PyObject *
3482compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3483{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003484 PyObject *seq1, *seq2;
3485 PyObject *data=NULL, *selectors=NULL;
3486 compressobject *lz;
3487 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003489 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3490 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003491
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003492 data = PyObject_GetIter(seq1);
3493 if (data == NULL)
3494 goto fail;
3495 selectors = PyObject_GetIter(seq2);
3496 if (selectors == NULL)
3497 goto fail;
3498
3499 /* create compressobject structure */
3500 lz = (compressobject *)type->tp_alloc(type, 0);
3501 if (lz == NULL)
3502 goto fail;
3503 lz->data = data;
3504 lz->selectors = selectors;
3505 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003506
3507fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003508 Py_XDECREF(data);
3509 Py_XDECREF(selectors);
3510 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003511}
3512
3513static void
3514compress_dealloc(compressobject *lz)
3515{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003516 PyObject_GC_UnTrack(lz);
3517 Py_XDECREF(lz->data);
3518 Py_XDECREF(lz->selectors);
3519 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003520}
3521
3522static int
3523compress_traverse(compressobject *lz, visitproc visit, void *arg)
3524{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003525 Py_VISIT(lz->data);
3526 Py_VISIT(lz->selectors);
3527 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003528}
3529
3530static PyObject *
3531compress_next(compressobject *lz)
3532{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003533 PyObject *data = lz->data, *selectors = lz->selectors;
3534 PyObject *datum, *selector;
3535 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3536 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3537 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003538
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003539 while (1) {
3540 /* Steps: get datum, get selector, evaluate selector.
3541 Order is important (to match the pure python version
3542 in terms of which input gets a chance to raise an
3543 exception first).
3544 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003545
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003546 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003547 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003548 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003549
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003550 selector = selectornext(selectors);
3551 if (selector == NULL) {
3552 Py_DECREF(datum);
3553 return NULL;
3554 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003555
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003556 ok = PyObject_IsTrue(selector);
3557 Py_DECREF(selector);
3558 if (ok == 1)
3559 return datum;
3560 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003561 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003562 return NULL;
3563 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003564}
3565
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003566static PyObject *
3567compress_reduce(compressobject *lz)
3568{
3569 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3570 lz->data, lz->selectors);
3571 }
3572
3573static PyMethodDef compress_methods[] = {
3574 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3575 reduce_doc},
3576 {NULL, NULL} /* sentinel */
3577};
3578
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003579PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003580"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003581\n\
3582Return data elements corresponding to true selector elements.\n\
3583Forms a shorter iterator from selected data elements using the\n\
3584selectors to choose the data elements.");
3585
3586static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003587 PyVarObject_HEAD_INIT(NULL, 0)
3588 "itertools.compress", /* tp_name */
3589 sizeof(compressobject), /* tp_basicsize */
3590 0, /* tp_itemsize */
3591 /* methods */
3592 (destructor)compress_dealloc, /* tp_dealloc */
3593 0, /* tp_print */
3594 0, /* tp_getattr */
3595 0, /* tp_setattr */
3596 0, /* tp_reserved */
3597 0, /* tp_repr */
3598 0, /* tp_as_number */
3599 0, /* tp_as_sequence */
3600 0, /* tp_as_mapping */
3601 0, /* tp_hash */
3602 0, /* tp_call */
3603 0, /* tp_str */
3604 PyObject_GenericGetAttr, /* tp_getattro */
3605 0, /* tp_setattro */
3606 0, /* tp_as_buffer */
3607 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3608 Py_TPFLAGS_BASETYPE, /* tp_flags */
3609 compress_doc, /* tp_doc */
3610 (traverseproc)compress_traverse, /* tp_traverse */
3611 0, /* tp_clear */
3612 0, /* tp_richcompare */
3613 0, /* tp_weaklistoffset */
3614 PyObject_SelfIter, /* tp_iter */
3615 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003616 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003617 0, /* tp_members */
3618 0, /* tp_getset */
3619 0, /* tp_base */
3620 0, /* tp_dict */
3621 0, /* tp_descr_get */
3622 0, /* tp_descr_set */
3623 0, /* tp_dictoffset */
3624 0, /* tp_init */
3625 0, /* tp_alloc */
3626 compress_new, /* tp_new */
3627 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003628};
3629
3630
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003631/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003632
3633typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003634 PyObject_HEAD
3635 PyObject *func;
3636 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003637} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003638
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003639static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003640
3641static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003642filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003643{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003644 PyObject *func, *seq;
3645 PyObject *it;
3646 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003647
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003648 if (type == &filterfalse_type &&
3649 !_PyArg_NoKeywords("filterfalse()", kwds))
3650 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003651
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003652 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3653 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003654
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003655 /* Get iterator. */
3656 it = PyObject_GetIter(seq);
3657 if (it == NULL)
3658 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003659
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003660 /* create filterfalseobject structure */
3661 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3662 if (lz == NULL) {
3663 Py_DECREF(it);
3664 return NULL;
3665 }
3666 Py_INCREF(func);
3667 lz->func = func;
3668 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003669
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003670 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003671}
3672
3673static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003674filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003675{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003676 PyObject_GC_UnTrack(lz);
3677 Py_XDECREF(lz->func);
3678 Py_XDECREF(lz->it);
3679 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003680}
3681
3682static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003683filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003685 Py_VISIT(lz->it);
3686 Py_VISIT(lz->func);
3687 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003688}
3689
3690static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003691filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003692{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003693 PyObject *item;
3694 PyObject *it = lz->it;
3695 long ok;
3696 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003698 iternext = *Py_TYPE(it)->tp_iternext;
3699 for (;;) {
3700 item = iternext(it);
3701 if (item == NULL)
3702 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003703
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3705 ok = PyObject_IsTrue(item);
3706 } else {
3707 PyObject *good;
3708 good = PyObject_CallFunctionObjArgs(lz->func,
3709 item, NULL);
3710 if (good == NULL) {
3711 Py_DECREF(item);
3712 return NULL;
3713 }
3714 ok = PyObject_IsTrue(good);
3715 Py_DECREF(good);
3716 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003717 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003718 return item;
3719 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003720 if (ok < 0)
3721 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003722 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003723}
3724
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003725static PyObject *
3726filterfalse_reduce(filterfalseobject *lz)
3727{
3728 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3729 lz->func, lz->it);
3730 }
3731
3732static PyMethodDef filterfalse_methods[] = {
3733 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3734 reduce_doc},
3735 {NULL, NULL} /* sentinel */
3736};
3737
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003738PyDoc_STRVAR(filterfalse_doc,
3739"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003740\n\
3741Return those items of sequence for which function(item) is false.\n\
3742If function is None, return the items that are false.");
3743
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003744static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003745 PyVarObject_HEAD_INIT(NULL, 0)
3746 "itertools.filterfalse", /* tp_name */
3747 sizeof(filterfalseobject), /* tp_basicsize */
3748 0, /* tp_itemsize */
3749 /* methods */
3750 (destructor)filterfalse_dealloc, /* tp_dealloc */
3751 0, /* tp_print */
3752 0, /* tp_getattr */
3753 0, /* tp_setattr */
3754 0, /* tp_reserved */
3755 0, /* tp_repr */
3756 0, /* tp_as_number */
3757 0, /* tp_as_sequence */
3758 0, /* tp_as_mapping */
3759 0, /* tp_hash */
3760 0, /* tp_call */
3761 0, /* tp_str */
3762 PyObject_GenericGetAttr, /* tp_getattro */
3763 0, /* tp_setattro */
3764 0, /* tp_as_buffer */
3765 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3766 Py_TPFLAGS_BASETYPE, /* tp_flags */
3767 filterfalse_doc, /* tp_doc */
3768 (traverseproc)filterfalse_traverse, /* tp_traverse */
3769 0, /* tp_clear */
3770 0, /* tp_richcompare */
3771 0, /* tp_weaklistoffset */
3772 PyObject_SelfIter, /* tp_iter */
3773 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003774 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003775 0, /* tp_members */
3776 0, /* tp_getset */
3777 0, /* tp_base */
3778 0, /* tp_dict */
3779 0, /* tp_descr_get */
3780 0, /* tp_descr_set */
3781 0, /* tp_dictoffset */
3782 0, /* tp_init */
3783 0, /* tp_alloc */
3784 filterfalse_new, /* tp_new */
3785 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003786};
3787
3788
3789/* count object ************************************************************/
3790
3791typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003792 PyObject_HEAD
3793 Py_ssize_t cnt;
3794 PyObject *long_cnt;
3795 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003796} countobject;
3797
Raymond Hettinger30729212009-02-12 06:28:27 +00003798/* Counting logic and invariants:
3799
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003800fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003801
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003802 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3803 Advances with: cnt += 1
3804 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003805
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003806slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003807
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003808 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3809 All counting is done with python objects (no overflows or underflows).
3810 Advances with: long_cnt += long_step
3811 Step may be zero -- effectively a slow version of repeat(cnt).
3812 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003813*/
3814
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003815static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003816
3817static PyObject *
3818count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3819{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003820 countobject *lz;
3821 int slow_mode = 0;
3822 Py_ssize_t cnt = 0;
3823 PyObject *long_cnt = NULL;
3824 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003825 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003826 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003827
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003828 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3829 kwlist, &long_cnt, &long_step))
3830 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003831
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003832 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3833 (long_step != NULL && !PyNumber_Check(long_step))) {
3834 PyErr_SetString(PyExc_TypeError, "a number is required");
3835 return NULL;
3836 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003837
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003838 if (long_cnt != NULL) {
3839 cnt = PyLong_AsSsize_t(long_cnt);
3840 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3841 PyErr_Clear();
3842 slow_mode = 1;
3843 }
3844 Py_INCREF(long_cnt);
3845 } else {
3846 cnt = 0;
3847 long_cnt = PyLong_FromLong(0);
3848 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003850 /* If not specified, step defaults to 1 */
3851 if (long_step == NULL) {
3852 long_step = PyLong_FromLong(1);
3853 if (long_step == NULL) {
3854 Py_DECREF(long_cnt);
3855 return NULL;
3856 }
3857 } else
3858 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003859
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003860 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003861
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003862 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003863 step = PyLong_AsLong(long_step);
3864 if (step != 1) {
3865 slow_mode = 1;
3866 if (step == -1 && PyErr_Occurred())
3867 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003868 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003869
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003870 if (slow_mode)
3871 cnt = PY_SSIZE_T_MAX;
3872 else
3873 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003875 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3876 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3877 assert(slow_mode ||
3878 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003880 /* create countobject structure */
3881 lz = (countobject *)type->tp_alloc(type, 0);
3882 if (lz == NULL) {
3883 Py_XDECREF(long_cnt);
3884 return NULL;
3885 }
3886 lz->cnt = cnt;
3887 lz->long_cnt = long_cnt;
3888 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003889
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003890 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003891}
3892
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003893static void
3894count_dealloc(countobject *lz)
3895{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003896 PyObject_GC_UnTrack(lz);
3897 Py_XDECREF(lz->long_cnt);
3898 Py_XDECREF(lz->long_step);
3899 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003900}
3901
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003902static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003903count_traverse(countobject *lz, visitproc visit, void *arg)
3904{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003905 Py_VISIT(lz->long_cnt);
3906 Py_VISIT(lz->long_step);
3907 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003908}
3909
3910static PyObject *
3911count_nextlong(countobject *lz)
3912{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 PyObject *long_cnt;
3914 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003915
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003916 long_cnt = lz->long_cnt;
3917 if (long_cnt == NULL) {
3918 /* Switch to slow_mode */
3919 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3920 if (long_cnt == NULL)
3921 return NULL;
3922 }
3923 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003925 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3926 if (stepped_up == NULL)
3927 return NULL;
3928 lz->long_cnt = stepped_up;
3929 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003930}
3931
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003932static PyObject *
3933count_next(countobject *lz)
3934{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 if (lz->cnt == PY_SSIZE_T_MAX)
3936 return count_nextlong(lz);
3937 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003938}
3939
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003940static PyObject *
3941count_repr(countobject *lz)
3942{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003943 if (lz->cnt != PY_SSIZE_T_MAX)
3944 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003945
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003946 if (PyLong_Check(lz->long_step)) {
3947 long step = PyLong_AsLong(lz->long_step);
3948 if (step == -1 && PyErr_Occurred()) {
3949 PyErr_Clear();
3950 }
3951 if (step == 1) {
3952 /* Don't display step when it is an integer equal to 1 */
3953 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3954 }
3955 }
3956 return PyUnicode_FromFormat("count(%R, %R)",
3957 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003958}
3959
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003960static PyObject *
3961count_reduce(countobject *lz)
3962{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003963 if (lz->cnt == PY_SSIZE_T_MAX)
3964 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3965 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003966}
3967
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003968static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003969 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003970 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003971 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003972};
3973
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003974PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003975 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003976\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003977Return a count object whose .__next__() method returns consecutive values.\n\
3978Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003979 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003980 x = firstval\n\
3981 while 1:\n\
3982 yield x\n\
3983 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003984
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003985static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003986 PyVarObject_HEAD_INIT(NULL, 0)
3987 "itertools.count", /* tp_name */
3988 sizeof(countobject), /* tp_basicsize */
3989 0, /* tp_itemsize */
3990 /* methods */
3991 (destructor)count_dealloc, /* tp_dealloc */
3992 0, /* tp_print */
3993 0, /* tp_getattr */
3994 0, /* tp_setattr */
3995 0, /* tp_reserved */
3996 (reprfunc)count_repr, /* tp_repr */
3997 0, /* tp_as_number */
3998 0, /* tp_as_sequence */
3999 0, /* tp_as_mapping */
4000 0, /* tp_hash */
4001 0, /* tp_call */
4002 0, /* tp_str */
4003 PyObject_GenericGetAttr, /* tp_getattro */
4004 0, /* tp_setattro */
4005 0, /* tp_as_buffer */
4006 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4007 Py_TPFLAGS_BASETYPE, /* tp_flags */
4008 count_doc, /* tp_doc */
4009 (traverseproc)count_traverse, /* tp_traverse */
4010 0, /* tp_clear */
4011 0, /* tp_richcompare */
4012 0, /* tp_weaklistoffset */
4013 PyObject_SelfIter, /* tp_iter */
4014 (iternextfunc)count_next, /* tp_iternext */
4015 count_methods, /* tp_methods */
4016 0, /* tp_members */
4017 0, /* tp_getset */
4018 0, /* tp_base */
4019 0, /* tp_dict */
4020 0, /* tp_descr_get */
4021 0, /* tp_descr_set */
4022 0, /* tp_dictoffset */
4023 0, /* tp_init */
4024 0, /* tp_alloc */
4025 count_new, /* tp_new */
4026 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004027};
4028
4029
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004030/* repeat object ************************************************************/
4031
4032typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004033 PyObject_HEAD
4034 PyObject *element;
4035 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004036} repeatobject;
4037
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004038static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004039
4040static PyObject *
4041repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4042{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 repeatobject *ro;
4044 PyObject *element;
4045 Py_ssize_t cnt = -1;
4046 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004047
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004048 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4049 &element, &cnt))
4050 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004051
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004052 if (PyTuple_Size(args) == 2 && cnt < 0)
4053 cnt = 0;
4054
4055 ro = (repeatobject *)type->tp_alloc(type, 0);
4056 if (ro == NULL)
4057 return NULL;
4058 Py_INCREF(element);
4059 ro->element = element;
4060 ro->cnt = cnt;
4061 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004062}
4063
4064static void
4065repeat_dealloc(repeatobject *ro)
4066{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004067 PyObject_GC_UnTrack(ro);
4068 Py_XDECREF(ro->element);
4069 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004070}
4071
4072static int
4073repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4074{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004075 Py_VISIT(ro->element);
4076 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004077}
4078
4079static PyObject *
4080repeat_next(repeatobject *ro)
4081{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004082 if (ro->cnt == 0)
4083 return NULL;
4084 if (ro->cnt > 0)
4085 ro->cnt--;
4086 Py_INCREF(ro->element);
4087 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004088}
4089
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004090static PyObject *
4091repeat_repr(repeatobject *ro)
4092{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004093 if (ro->cnt == -1)
4094 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4095 else
4096 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4097}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004098
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004099static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004100repeat_len(repeatobject *ro)
4101{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004102 if (ro->cnt == -1) {
4103 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4104 return NULL;
4105 }
4106 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004107}
4108
Armin Rigof5b3e362006-02-11 21:32:43 +00004109PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004110
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004111static PyObject *
4112repeat_reduce(repeatobject *ro)
4113{
4114 /* unpickle this so that a new repeat iterator is constructed with an
4115 * object, then call __setstate__ on it to set cnt
4116 */
4117 if (ro->cnt >= 0)
4118 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4119 else
4120 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4121}
4122
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004123static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004124 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004125 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004126 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004127};
4128
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004129PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004130"repeat(object [,times]) -> create an iterator which returns the object\n\
4131for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004132endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004133
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004134static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004135 PyVarObject_HEAD_INIT(NULL, 0)
4136 "itertools.repeat", /* tp_name */
4137 sizeof(repeatobject), /* tp_basicsize */
4138 0, /* tp_itemsize */
4139 /* methods */
4140 (destructor)repeat_dealloc, /* tp_dealloc */
4141 0, /* tp_print */
4142 0, /* tp_getattr */
4143 0, /* tp_setattr */
4144 0, /* tp_reserved */
4145 (reprfunc)repeat_repr, /* tp_repr */
4146 0, /* tp_as_number */
4147 0, /* tp_as_sequence */
4148 0, /* tp_as_mapping */
4149 0, /* tp_hash */
4150 0, /* tp_call */
4151 0, /* tp_str */
4152 PyObject_GenericGetAttr, /* tp_getattro */
4153 0, /* tp_setattro */
4154 0, /* tp_as_buffer */
4155 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4156 Py_TPFLAGS_BASETYPE, /* tp_flags */
4157 repeat_doc, /* tp_doc */
4158 (traverseproc)repeat_traverse, /* tp_traverse */
4159 0, /* tp_clear */
4160 0, /* tp_richcompare */
4161 0, /* tp_weaklistoffset */
4162 PyObject_SelfIter, /* tp_iter */
4163 (iternextfunc)repeat_next, /* tp_iternext */
4164 repeat_methods, /* tp_methods */
4165 0, /* tp_members */
4166 0, /* tp_getset */
4167 0, /* tp_base */
4168 0, /* tp_dict */
4169 0, /* tp_descr_get */
4170 0, /* tp_descr_set */
4171 0, /* tp_dictoffset */
4172 0, /* tp_init */
4173 0, /* tp_alloc */
4174 repeat_new, /* tp_new */
4175 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004176};
4177
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004178/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004179
4180#include "Python.h"
4181
4182typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004183 PyObject_HEAD
4184 Py_ssize_t tuplesize;
4185 Py_ssize_t numactive;
4186 PyObject *ittuple; /* tuple of iterators */
4187 PyObject *result;
4188 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004189} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004190
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004191static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004192
4193static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004194zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004195{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 ziplongestobject *lz;
4197 Py_ssize_t i;
4198 PyObject *ittuple; /* tuple of iterators */
4199 PyObject *result;
4200 PyObject *fillvalue = Py_None;
4201 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004202
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004203 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4204 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4205 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4206 PyErr_SetString(PyExc_TypeError,
4207 "zip_longest() got an unexpected keyword argument");
4208 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004209 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004210 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004211
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004212 /* args must be a tuple */
4213 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004214
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004215 /* obtain iterators */
4216 ittuple = PyTuple_New(tuplesize);
4217 if (ittuple == NULL)
4218 return NULL;
4219 for (i=0; i < tuplesize; ++i) {
4220 PyObject *item = PyTuple_GET_ITEM(args, i);
4221 PyObject *it = PyObject_GetIter(item);
4222 if (it == NULL) {
4223 if (PyErr_ExceptionMatches(PyExc_TypeError))
4224 PyErr_Format(PyExc_TypeError,
4225 "zip_longest argument #%zd must support iteration",
4226 i+1);
4227 Py_DECREF(ittuple);
4228 return NULL;
4229 }
4230 PyTuple_SET_ITEM(ittuple, i, it);
4231 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004232
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004233 /* create a result holder */
4234 result = PyTuple_New(tuplesize);
4235 if (result == NULL) {
4236 Py_DECREF(ittuple);
4237 return NULL;
4238 }
4239 for (i=0 ; i < tuplesize ; i++) {
4240 Py_INCREF(Py_None);
4241 PyTuple_SET_ITEM(result, i, Py_None);
4242 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004243
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004244 /* create ziplongestobject structure */
4245 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4246 if (lz == NULL) {
4247 Py_DECREF(ittuple);
4248 Py_DECREF(result);
4249 return NULL;
4250 }
4251 lz->ittuple = ittuple;
4252 lz->tuplesize = tuplesize;
4253 lz->numactive = tuplesize;
4254 lz->result = result;
4255 Py_INCREF(fillvalue);
4256 lz->fillvalue = fillvalue;
4257 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004258}
4259
4260static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004261zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004262{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004263 PyObject_GC_UnTrack(lz);
4264 Py_XDECREF(lz->ittuple);
4265 Py_XDECREF(lz->result);
4266 Py_XDECREF(lz->fillvalue);
4267 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004268}
4269
4270static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004271zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004272{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004273 Py_VISIT(lz->ittuple);
4274 Py_VISIT(lz->result);
4275 Py_VISIT(lz->fillvalue);
4276 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004277}
4278
4279static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004280zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004281{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004282 Py_ssize_t i;
4283 Py_ssize_t tuplesize = lz->tuplesize;
4284 PyObject *result = lz->result;
4285 PyObject *it;
4286 PyObject *item;
4287 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004288
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004289 if (tuplesize == 0)
4290 return NULL;
4291 if (lz->numactive == 0)
4292 return NULL;
4293 if (Py_REFCNT(result) == 1) {
4294 Py_INCREF(result);
4295 for (i=0 ; i < tuplesize ; i++) {
4296 it = PyTuple_GET_ITEM(lz->ittuple, i);
4297 if (it == NULL) {
4298 Py_INCREF(lz->fillvalue);
4299 item = lz->fillvalue;
4300 } else {
4301 item = PyIter_Next(it);
4302 if (item == NULL) {
4303 lz->numactive -= 1;
4304 if (lz->numactive == 0 || PyErr_Occurred()) {
4305 lz->numactive = 0;
4306 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004307 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004308 } else {
4309 Py_INCREF(lz->fillvalue);
4310 item = lz->fillvalue;
4311 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4312 Py_DECREF(it);
4313 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004314 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004315 }
4316 olditem = PyTuple_GET_ITEM(result, i);
4317 PyTuple_SET_ITEM(result, i, item);
4318 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004319 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004320 } else {
4321 result = PyTuple_New(tuplesize);
4322 if (result == NULL)
4323 return NULL;
4324 for (i=0 ; i < tuplesize ; i++) {
4325 it = PyTuple_GET_ITEM(lz->ittuple, i);
4326 if (it == NULL) {
4327 Py_INCREF(lz->fillvalue);
4328 item = lz->fillvalue;
4329 } else {
4330 item = PyIter_Next(it);
4331 if (item == NULL) {
4332 lz->numactive -= 1;
4333 if (lz->numactive == 0 || PyErr_Occurred()) {
4334 lz->numactive = 0;
4335 Py_DECREF(result);
4336 return NULL;
4337 } else {
4338 Py_INCREF(lz->fillvalue);
4339 item = lz->fillvalue;
4340 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4341 Py_DECREF(it);
4342 }
4343 }
4344 }
4345 PyTuple_SET_ITEM(result, i, item);
4346 }
4347 }
4348 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004349}
4350
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004351static PyObject *
4352zip_longest_reduce(ziplongestobject *lz)
4353{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004354
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004355 /* Create a new tuple with empty sequences where appropriate to pickle.
4356 * Then use setstate to set the fillvalue
4357 */
4358 int i;
4359 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4360 if (args == NULL)
4361 return NULL;
4362 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4363 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4364 if (elem == NULL) {
4365 elem = PyTuple_New(0);
4366 if (elem == NULL) {
4367 Py_DECREF(args);
4368 return NULL;
4369 }
4370 } else
4371 Py_INCREF(elem);
4372 PyTuple_SET_ITEM(args, i, elem);
4373 }
4374 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4375}
4376
4377static PyObject *
4378zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4379{
4380 Py_CLEAR(lz->fillvalue);
4381 lz->fillvalue = state;
4382 Py_INCREF(lz->fillvalue);
4383 Py_RETURN_NONE;
4384}
4385
4386static PyMethodDef zip_longest_methods[] = {
4387 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4388 reduce_doc},
4389 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4390 setstate_doc},
4391 {NULL, NULL} /* sentinel */
4392};
4393
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004394PyDoc_STRVAR(zip_longest_doc,
4395"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004396\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004397Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004398the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004399method continues until the longest iterable in the argument sequence\n\
4400is exhausted and then it raises StopIteration. When the shorter iterables\n\
4401are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4402defaults to None or can be specified by a keyword argument.\n\
4403");
4404
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004405static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004406 PyVarObject_HEAD_INIT(NULL, 0)
4407 "itertools.zip_longest", /* tp_name */
4408 sizeof(ziplongestobject), /* tp_basicsize */
4409 0, /* tp_itemsize */
4410 /* methods */
4411 (destructor)zip_longest_dealloc, /* tp_dealloc */
4412 0, /* tp_print */
4413 0, /* tp_getattr */
4414 0, /* tp_setattr */
4415 0, /* tp_reserved */
4416 0, /* tp_repr */
4417 0, /* tp_as_number */
4418 0, /* tp_as_sequence */
4419 0, /* tp_as_mapping */
4420 0, /* tp_hash */
4421 0, /* tp_call */
4422 0, /* tp_str */
4423 PyObject_GenericGetAttr, /* tp_getattro */
4424 0, /* tp_setattro */
4425 0, /* tp_as_buffer */
4426 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4427 Py_TPFLAGS_BASETYPE, /* tp_flags */
4428 zip_longest_doc, /* tp_doc */
4429 (traverseproc)zip_longest_traverse, /* tp_traverse */
4430 0, /* tp_clear */
4431 0, /* tp_richcompare */
4432 0, /* tp_weaklistoffset */
4433 PyObject_SelfIter, /* tp_iter */
4434 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004435 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004436 0, /* tp_members */
4437 0, /* tp_getset */
4438 0, /* tp_base */
4439 0, /* tp_dict */
4440 0, /* tp_descr_get */
4441 0, /* tp_descr_set */
4442 0, /* tp_dictoffset */
4443 0, /* tp_init */
4444 0, /* tp_alloc */
4445 zip_longest_new, /* tp_new */
4446 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004447};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004448
4449/* module level code ********************************************************/
4450
4451PyDoc_STRVAR(module_doc,
4452"Functional tools for creating and using iterators.\n\
4453\n\
4454Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004455count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004456cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004457repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004458\n\
4459Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004460accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004461chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
4462compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4463dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4464groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004465filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004466islice(seq, [start,] stop [, step]) --> elements from\n\
4467 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004468starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004469tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004470takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004471zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004472\n\
4473Combinatoric generators:\n\
4474product(p, q, ... [repeat=1]) --> cartesian product\n\
4475permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004476combinations(p, r)\n\
4477combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004478");
4479
4480
Raymond Hettingerad983e72003-11-12 14:32:26 +00004481static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004482 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4483 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004484};
4485
Martin v. Löwis1a214512008-06-11 05:26:20 +00004486
4487static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004488 PyModuleDef_HEAD_INIT,
4489 "itertools",
4490 module_doc,
4491 -1,
4492 module_methods,
4493 NULL,
4494 NULL,
4495 NULL,
4496 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004497};
4498
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004499PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004500PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004501{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004502 int i;
4503 PyObject *m;
4504 char *name;
4505 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004506 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004507 &combinations_type,
4508 &cwr_type,
4509 &cycle_type,
4510 &dropwhile_type,
4511 &takewhile_type,
4512 &islice_type,
4513 &starmap_type,
4514 &chain_type,
4515 &compress_type,
4516 &filterfalse_type,
4517 &count_type,
4518 &ziplongest_type,
4519 &permutations_type,
4520 &product_type,
4521 &repeat_type,
4522 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004523 &_grouper_type,
4524 &tee_type,
4525 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004526 NULL
4527 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004528
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004529 Py_TYPE(&teedataobject_type) = &PyType_Type;
4530 m = PyModule_Create(&itertoolsmodule);
4531 if (m == NULL)
4532 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004533
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004534 for (i=0 ; typelist[i] != NULL ; i++) {
4535 if (PyType_Ready(typelist[i]) < 0)
4536 return NULL;
4537 name = strchr(typelist[i]->tp_name, '.');
4538 assert (name != NULL);
4539 Py_INCREF(typelist[i]);
4540 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4541 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004542
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004543 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004544}