blob: 9115b67146c00f9c1a6e0162b0245b54ea9ec52f [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
476static int
477teedataobject_clear(teedataobject *tdo)
478{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000479 int i;
480 Py_CLEAR(tdo->it);
481 for (i=0 ; i<tdo->numread ; i++)
482 Py_CLEAR(tdo->values[i]);
483 Py_CLEAR(tdo->nextlink);
484 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000485}
486
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000487static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000488teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 PyObject_GC_UnTrack(tdo);
491 teedataobject_clear(tdo);
492 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000493}
494
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000495static PyObject *
496teedataobject_reduce(teedataobject *tdo)
497{
498 int i;
499 /* create a temporary list of already iterated values */
500 PyObject *values = PyList_New(tdo->numread);
501 if (!values)
502 return NULL;
503 for (i=0 ; i<tdo->numread ; i++) {
504 Py_INCREF(tdo->values[i]);
505 PyList_SET_ITEM(values, i, tdo->values[i]);
506 }
507 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
508 values,
509 tdo->nextlink ? tdo->nextlink : Py_None);
510}
511
512static PyTypeObject teedataobject_type;
513
514static PyObject *
515teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
516{
517 teedataobject *tdo;
518 PyObject *it, *values, *next;
519 Py_ssize_t i, len;
520
521 assert(type == &teedataobject_type);
522 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
523 return NULL;
524
525 tdo = (teedataobject *)teedataobject_newinternal(it);
526 if (!tdo)
527 return NULL;
528
529 len = PyList_GET_SIZE(values);
530 if (len > LINKCELLS)
531 goto err;
532 for (i=0; i<len; i++) {
533 tdo->values[i] = PyList_GET_ITEM(values, i);
534 Py_INCREF(tdo->values[i]);
535 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200536 /* len <= LINKCELLS < INT_MAX */
537 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000538
539 if (len == LINKCELLS) {
540 if (next != Py_None) {
541 if (Py_TYPE(next) != &teedataobject_type)
542 goto err;
543 assert(tdo->nextlink == NULL);
544 Py_INCREF(next);
545 tdo->nextlink = next;
546 }
547 } else {
548 if (next != Py_None)
549 goto err; /* shouldn't have a next if we are not full */
550 }
551 return (PyObject*)tdo;
552
553err:
554 Py_XDECREF(tdo);
555 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
556 return NULL;
557}
558
559static PyMethodDef teedataobject_methods[] = {
560 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
561 reduce_doc},
562 {NULL, NULL} /* sentinel */
563};
564
Raymond Hettingerad983e72003-11-12 14:32:26 +0000565PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000566
Raymond Hettingerad983e72003-11-12 14:32:26 +0000567static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000568 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000569 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000570 sizeof(teedataobject), /* tp_basicsize */
571 0, /* tp_itemsize */
572 /* methods */
573 (destructor)teedataobject_dealloc, /* tp_dealloc */
574 0, /* tp_print */
575 0, /* tp_getattr */
576 0, /* tp_setattr */
577 0, /* tp_reserved */
578 0, /* tp_repr */
579 0, /* tp_as_number */
580 0, /* tp_as_sequence */
581 0, /* tp_as_mapping */
582 0, /* tp_hash */
583 0, /* tp_call */
584 0, /* tp_str */
585 PyObject_GenericGetAttr, /* tp_getattro */
586 0, /* tp_setattro */
587 0, /* tp_as_buffer */
588 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
589 teedataobject_doc, /* tp_doc */
590 (traverseproc)teedataobject_traverse, /* tp_traverse */
591 (inquiry)teedataobject_clear, /* tp_clear */
592 0, /* tp_richcompare */
593 0, /* tp_weaklistoffset */
594 0, /* tp_iter */
595 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000596 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000597 0, /* tp_members */
598 0, /* tp_getset */
599 0, /* tp_base */
600 0, /* tp_dict */
601 0, /* tp_descr_get */
602 0, /* tp_descr_set */
603 0, /* tp_dictoffset */
604 0, /* tp_init */
605 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000606 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000607 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000608};
609
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000610
611static PyTypeObject tee_type;
612
613static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000614tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000615{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000616 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000617
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000618 if (to->index >= LINKCELLS) {
619 link = teedataobject_jumplink(to->dataobj);
620 Py_DECREF(to->dataobj);
621 to->dataobj = (teedataobject *)link;
622 to->index = 0;
623 }
624 value = teedataobject_getitem(to->dataobj, to->index);
625 if (value == NULL)
626 return NULL;
627 to->index++;
628 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629}
630
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000631static int
632tee_traverse(teeobject *to, visitproc visit, void *arg)
633{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000634 Py_VISIT((PyObject *)to->dataobj);
635 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000636}
637
Raymond Hettingerad983e72003-11-12 14:32:26 +0000638static PyObject *
639tee_copy(teeobject *to)
640{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000641 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000643 newto = PyObject_GC_New(teeobject, &tee_type);
644 if (newto == NULL)
645 return NULL;
646 Py_INCREF(to->dataobj);
647 newto->dataobj = to->dataobj;
648 newto->index = to->index;
649 newto->weakreflist = NULL;
650 PyObject_GC_Track(newto);
651 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000652}
653
654PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
655
656static PyObject *
657tee_fromiterable(PyObject *iterable)
658{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000659 teeobject *to;
660 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 it = PyObject_GetIter(iterable);
663 if (it == NULL)
664 return NULL;
665 if (PyObject_TypeCheck(it, &tee_type)) {
666 to = (teeobject *)tee_copy((teeobject *)it);
667 goto done;
668 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000669
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000670 to = PyObject_GC_New(teeobject, &tee_type);
671 if (to == NULL)
672 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000673 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000674 if (!to->dataobj) {
675 PyObject_GC_Del(to);
676 to = NULL;
677 goto done;
678 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000679
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000680 to->index = 0;
681 to->weakreflist = NULL;
682 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000683done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000684 Py_XDECREF(it);
685 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000686}
687
688static PyObject *
689tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
690{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000691 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000692
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000693 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000694 return NULL;
695 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000696}
697
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000698static int
699tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000700{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000701 if (to->weakreflist != NULL)
702 PyObject_ClearWeakRefs((PyObject *) to);
703 Py_CLEAR(to->dataobj);
704 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000705}
706
707static void
708tee_dealloc(teeobject *to)
709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000710 PyObject_GC_UnTrack(to);
711 tee_clear(to);
712 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713}
714
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000715static PyObject *
716tee_reduce(teeobject *to)
717{
718 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
719}
720
721static PyObject *
722tee_setstate(teeobject *to, PyObject *state)
723{
724 teedataobject *tdo;
725 int index;
726 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
727 return NULL;
728 if (index < 0 || index > LINKCELLS) {
729 PyErr_SetString(PyExc_ValueError, "Index out of range");
730 return NULL;
731 }
732 Py_CLEAR(to->dataobj);
733 to->dataobj = tdo;
734 Py_INCREF(to->dataobj);
735 to->index = index;
736 Py_RETURN_NONE;
737}
738
Raymond Hettingerad983e72003-11-12 14:32:26 +0000739PyDoc_STRVAR(teeobject_doc,
740"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000741
Raymond Hettingerad983e72003-11-12 14:32:26 +0000742static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000743 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000744 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
745 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000746 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000747};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000748
749static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000750 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000751 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000752 sizeof(teeobject), /* tp_basicsize */
753 0, /* tp_itemsize */
754 /* methods */
755 (destructor)tee_dealloc, /* tp_dealloc */
756 0, /* tp_print */
757 0, /* tp_getattr */
758 0, /* tp_setattr */
759 0, /* tp_reserved */
760 0, /* tp_repr */
761 0, /* tp_as_number */
762 0, /* tp_as_sequence */
763 0, /* tp_as_mapping */
764 0, /* tp_hash */
765 0, /* tp_call */
766 0, /* tp_str */
767 0, /* tp_getattro */
768 0, /* tp_setattro */
769 0, /* tp_as_buffer */
770 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
771 teeobject_doc, /* tp_doc */
772 (traverseproc)tee_traverse, /* tp_traverse */
773 (inquiry)tee_clear, /* tp_clear */
774 0, /* tp_richcompare */
775 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
776 PyObject_SelfIter, /* tp_iter */
777 (iternextfunc)tee_next, /* tp_iternext */
778 tee_methods, /* tp_methods */
779 0, /* tp_members */
780 0, /* tp_getset */
781 0, /* tp_base */
782 0, /* tp_dict */
783 0, /* tp_descr_get */
784 0, /* tp_descr_set */
785 0, /* tp_dictoffset */
786 0, /* tp_init */
787 0, /* tp_alloc */
788 tee_new, /* tp_new */
789 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000790};
791
Raymond Hettingerad983e72003-11-12 14:32:26 +0000792static PyObject *
793tee(PyObject *self, PyObject *args)
794{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000795 Py_ssize_t i, n=2;
796 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200797 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000798
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000799 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
800 return NULL;
801 if (n < 0) {
802 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
803 return NULL;
804 }
805 result = PyTuple_New(n);
806 if (result == NULL)
807 return NULL;
808 if (n == 0)
809 return result;
810 it = PyObject_GetIter(iterable);
811 if (it == NULL) {
812 Py_DECREF(result);
813 return NULL;
814 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200815 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000816 copyable = tee_fromiterable(it);
817 Py_DECREF(it);
818 if (copyable == NULL) {
819 Py_DECREF(result);
820 return NULL;
821 }
822 } else
823 copyable = it;
824 PyTuple_SET_ITEM(result, 0, copyable);
825 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200826
827 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000828 if (copyable == NULL) {
829 Py_DECREF(result);
830 return NULL;
831 }
832 PyTuple_SET_ITEM(result, i, copyable);
833 }
834 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000835}
836
837PyDoc_STRVAR(tee_doc,
838"tee(iterable, n=2) --> tuple of n independent iterators.");
839
840
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000841/* cycle object **********************************************************/
842
843typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000844 PyObject_HEAD
845 PyObject *it;
846 PyObject *saved;
847 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000848} cycleobject;
849
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000850static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000851
852static PyObject *
853cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
854{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000855 PyObject *it;
856 PyObject *iterable;
857 PyObject *saved;
858 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000859
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000860 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
861 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000862
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
864 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000865
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000866 /* Get iterator. */
867 it = PyObject_GetIter(iterable);
868 if (it == NULL)
869 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000870
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000871 saved = PyList_New(0);
872 if (saved == NULL) {
873 Py_DECREF(it);
874 return NULL;
875 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000877 /* create cycleobject structure */
878 lz = (cycleobject *)type->tp_alloc(type, 0);
879 if (lz == NULL) {
880 Py_DECREF(it);
881 Py_DECREF(saved);
882 return NULL;
883 }
884 lz->it = it;
885 lz->saved = saved;
886 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000887
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000888 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000889}
890
891static void
892cycle_dealloc(cycleobject *lz)
893{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000894 PyObject_GC_UnTrack(lz);
895 Py_XDECREF(lz->saved);
896 Py_XDECREF(lz->it);
897 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000898}
899
900static int
901cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
902{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000903 Py_VISIT(lz->it);
904 Py_VISIT(lz->saved);
905 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906}
907
908static PyObject *
909cycle_next(cycleobject *lz)
910{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000911 PyObject *item;
912 PyObject *it;
913 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000914
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000915 while (1) {
916 item = PyIter_Next(lz->it);
917 if (item != NULL) {
918 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
919 Py_DECREF(item);
920 return NULL;
921 }
922 return item;
923 }
924 if (PyErr_Occurred()) {
925 if (PyErr_ExceptionMatches(PyExc_StopIteration))
926 PyErr_Clear();
927 else
928 return NULL;
929 }
930 if (PyList_Size(lz->saved) == 0)
931 return NULL;
932 it = PyObject_GetIter(lz->saved);
933 if (it == NULL)
934 return NULL;
935 tmp = lz->it;
936 lz->it = it;
937 lz->firstpass = 1;
938 Py_DECREF(tmp);
939 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000940}
941
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000942static PyObject *
943cycle_reduce(cycleobject *lz)
944{
945 /* Create a new cycle with the iterator tuple, then set
946 * the saved state on it.
947 */
948 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
949 lz->it, lz->saved, lz->firstpass);
950 }
951
952static PyObject *
953cycle_setstate(cycleobject *lz, PyObject *state)
954{
955 PyObject *saved=NULL;
956 int firstpass;
957 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
958 return NULL;
959 Py_CLEAR(lz->saved);
960 lz->saved = saved;
961 Py_XINCREF(lz->saved);
962 lz->firstpass = firstpass != 0;
963 Py_RETURN_NONE;
964}
965
966static PyMethodDef cycle_methods[] = {
967 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
968 reduce_doc},
969 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
970 setstate_doc},
971 {NULL, NULL} /* sentinel */
972};
973
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000974PyDoc_STRVAR(cycle_doc,
975"cycle(iterable) --> cycle object\n\
976\n\
977Return elements from the iterable until it is exhausted.\n\
978Then repeat the sequence indefinitely.");
979
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000980static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000981 PyVarObject_HEAD_INIT(NULL, 0)
982 "itertools.cycle", /* tp_name */
983 sizeof(cycleobject), /* tp_basicsize */
984 0, /* tp_itemsize */
985 /* methods */
986 (destructor)cycle_dealloc, /* tp_dealloc */
987 0, /* tp_print */
988 0, /* tp_getattr */
989 0, /* tp_setattr */
990 0, /* tp_reserved */
991 0, /* tp_repr */
992 0, /* tp_as_number */
993 0, /* tp_as_sequence */
994 0, /* tp_as_mapping */
995 0, /* tp_hash */
996 0, /* tp_call */
997 0, /* tp_str */
998 PyObject_GenericGetAttr, /* tp_getattro */
999 0, /* tp_setattro */
1000 0, /* tp_as_buffer */
1001 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1002 Py_TPFLAGS_BASETYPE, /* tp_flags */
1003 cycle_doc, /* tp_doc */
1004 (traverseproc)cycle_traverse, /* tp_traverse */
1005 0, /* tp_clear */
1006 0, /* tp_richcompare */
1007 0, /* tp_weaklistoffset */
1008 PyObject_SelfIter, /* tp_iter */
1009 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001010 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001011 0, /* tp_members */
1012 0, /* tp_getset */
1013 0, /* tp_base */
1014 0, /* tp_dict */
1015 0, /* tp_descr_get */
1016 0, /* tp_descr_set */
1017 0, /* tp_dictoffset */
1018 0, /* tp_init */
1019 0, /* tp_alloc */
1020 cycle_new, /* tp_new */
1021 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001022};
1023
1024
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025/* dropwhile object **********************************************************/
1026
1027typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001028 PyObject_HEAD
1029 PyObject *func;
1030 PyObject *it;
1031 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001032} dropwhileobject;
1033
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001034static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001035
1036static PyObject *
1037dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1038{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001039 PyObject *func, *seq;
1040 PyObject *it;
1041 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001042
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001043 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1044 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001046 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1047 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001048
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001049 /* Get iterator. */
1050 it = PyObject_GetIter(seq);
1051 if (it == NULL)
1052 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001054 /* create dropwhileobject structure */
1055 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1056 if (lz == NULL) {
1057 Py_DECREF(it);
1058 return NULL;
1059 }
1060 Py_INCREF(func);
1061 lz->func = func;
1062 lz->it = it;
1063 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066}
1067
1068static void
1069dropwhile_dealloc(dropwhileobject *lz)
1070{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001071 PyObject_GC_UnTrack(lz);
1072 Py_XDECREF(lz->func);
1073 Py_XDECREF(lz->it);
1074 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075}
1076
1077static int
1078dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1079{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001080 Py_VISIT(lz->it);
1081 Py_VISIT(lz->func);
1082 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083}
1084
1085static PyObject *
1086dropwhile_next(dropwhileobject *lz)
1087{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001088 PyObject *item, *good;
1089 PyObject *it = lz->it;
1090 long ok;
1091 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001093 iternext = *Py_TYPE(it)->tp_iternext;
1094 for (;;) {
1095 item = iternext(it);
1096 if (item == NULL)
1097 return NULL;
1098 if (lz->start == 1)
1099 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001101 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1102 if (good == NULL) {
1103 Py_DECREF(item);
1104 return NULL;
1105 }
1106 ok = PyObject_IsTrue(good);
1107 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001108 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001109 lz->start = 1;
1110 return item;
1111 }
1112 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001113 if (ok < 0)
1114 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001115 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116}
1117
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001118static PyObject *
1119dropwhile_reduce(dropwhileobject *lz)
1120{
1121 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1122 lz->func, lz->it, lz->start);
1123}
1124
1125static PyObject *
1126dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1127{
1128 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001129 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001130 return NULL;
1131 lz->start = start;
1132 Py_RETURN_NONE;
1133}
1134
1135static PyMethodDef dropwhile_methods[] = {
1136 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1137 reduce_doc},
1138 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1139 setstate_doc},
1140 {NULL, NULL} /* sentinel */
1141};
1142
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001143PyDoc_STRVAR(dropwhile_doc,
1144"dropwhile(predicate, iterable) --> dropwhile object\n\
1145\n\
1146Drop items from the iterable while predicate(item) is true.\n\
1147Afterwards, return every element until the iterable is exhausted.");
1148
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001149static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001150 PyVarObject_HEAD_INIT(NULL, 0)
1151 "itertools.dropwhile", /* tp_name */
1152 sizeof(dropwhileobject), /* tp_basicsize */
1153 0, /* tp_itemsize */
1154 /* methods */
1155 (destructor)dropwhile_dealloc, /* tp_dealloc */
1156 0, /* tp_print */
1157 0, /* tp_getattr */
1158 0, /* tp_setattr */
1159 0, /* tp_reserved */
1160 0, /* tp_repr */
1161 0, /* tp_as_number */
1162 0, /* tp_as_sequence */
1163 0, /* tp_as_mapping */
1164 0, /* tp_hash */
1165 0, /* tp_call */
1166 0, /* tp_str */
1167 PyObject_GenericGetAttr, /* tp_getattro */
1168 0, /* tp_setattro */
1169 0, /* tp_as_buffer */
1170 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1171 Py_TPFLAGS_BASETYPE, /* tp_flags */
1172 dropwhile_doc, /* tp_doc */
1173 (traverseproc)dropwhile_traverse, /* tp_traverse */
1174 0, /* tp_clear */
1175 0, /* tp_richcompare */
1176 0, /* tp_weaklistoffset */
1177 PyObject_SelfIter, /* tp_iter */
1178 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001179 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001180 0, /* tp_members */
1181 0, /* tp_getset */
1182 0, /* tp_base */
1183 0, /* tp_dict */
1184 0, /* tp_descr_get */
1185 0, /* tp_descr_set */
1186 0, /* tp_dictoffset */
1187 0, /* tp_init */
1188 0, /* tp_alloc */
1189 dropwhile_new, /* tp_new */
1190 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001191};
1192
1193
1194/* takewhile object **********************************************************/
1195
1196typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001197 PyObject_HEAD
1198 PyObject *func;
1199 PyObject *it;
1200 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001201} takewhileobject;
1202
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001203static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204
1205static PyObject *
1206takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1207{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001208 PyObject *func, *seq;
1209 PyObject *it;
1210 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001212 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1213 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001214
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001215 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1216 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001218 /* Get iterator. */
1219 it = PyObject_GetIter(seq);
1220 if (it == NULL)
1221 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001222
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001223 /* create takewhileobject structure */
1224 lz = (takewhileobject *)type->tp_alloc(type, 0);
1225 if (lz == NULL) {
1226 Py_DECREF(it);
1227 return NULL;
1228 }
1229 Py_INCREF(func);
1230 lz->func = func;
1231 lz->it = it;
1232 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001235}
1236
1237static void
1238takewhile_dealloc(takewhileobject *lz)
1239{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001240 PyObject_GC_UnTrack(lz);
1241 Py_XDECREF(lz->func);
1242 Py_XDECREF(lz->it);
1243 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244}
1245
1246static int
1247takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1248{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001249 Py_VISIT(lz->it);
1250 Py_VISIT(lz->func);
1251 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001252}
1253
1254static PyObject *
1255takewhile_next(takewhileobject *lz)
1256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001257 PyObject *item, *good;
1258 PyObject *it = lz->it;
1259 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001260
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001261 if (lz->stop == 1)
1262 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001263
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001264 item = (*Py_TYPE(it)->tp_iternext)(it);
1265 if (item == NULL)
1266 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001267
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001268 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1269 if (good == NULL) {
1270 Py_DECREF(item);
1271 return NULL;
1272 }
1273 ok = PyObject_IsTrue(good);
1274 Py_DECREF(good);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001275 if (ok == 1)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001276 return item;
1277 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001278 if (ok == 0)
1279 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001280 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001281}
1282
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001283static PyObject *
1284takewhile_reduce(takewhileobject *lz)
1285{
1286 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1287 lz->func, lz->it, lz->stop);
1288}
1289
1290static PyObject *
1291takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1292{
1293 int stop = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001294 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001295 return NULL;
1296 lz->stop = stop;
1297 Py_RETURN_NONE;
1298}
1299
1300static PyMethodDef takewhile_reduce_methods[] = {
1301 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1302 reduce_doc},
1303 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1304 setstate_doc},
1305 {NULL, NULL} /* sentinel */
1306};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307PyDoc_STRVAR(takewhile_doc,
1308"takewhile(predicate, iterable) --> takewhile object\n\
1309\n\
1310Return successive entries from an iterable as long as the \n\
1311predicate evaluates to true for each entry.");
1312
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001313static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001314 PyVarObject_HEAD_INIT(NULL, 0)
1315 "itertools.takewhile", /* tp_name */
1316 sizeof(takewhileobject), /* tp_basicsize */
1317 0, /* tp_itemsize */
1318 /* methods */
1319 (destructor)takewhile_dealloc, /* tp_dealloc */
1320 0, /* tp_print */
1321 0, /* tp_getattr */
1322 0, /* tp_setattr */
1323 0, /* tp_reserved */
1324 0, /* tp_repr */
1325 0, /* tp_as_number */
1326 0, /* tp_as_sequence */
1327 0, /* tp_as_mapping */
1328 0, /* tp_hash */
1329 0, /* tp_call */
1330 0, /* tp_str */
1331 PyObject_GenericGetAttr, /* tp_getattro */
1332 0, /* tp_setattro */
1333 0, /* tp_as_buffer */
1334 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1335 Py_TPFLAGS_BASETYPE, /* tp_flags */
1336 takewhile_doc, /* tp_doc */
1337 (traverseproc)takewhile_traverse, /* tp_traverse */
1338 0, /* tp_clear */
1339 0, /* tp_richcompare */
1340 0, /* tp_weaklistoffset */
1341 PyObject_SelfIter, /* tp_iter */
1342 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001343 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001344 0, /* tp_members */
1345 0, /* tp_getset */
1346 0, /* tp_base */
1347 0, /* tp_dict */
1348 0, /* tp_descr_get */
1349 0, /* tp_descr_set */
1350 0, /* tp_dictoffset */
1351 0, /* tp_init */
1352 0, /* tp_alloc */
1353 takewhile_new, /* tp_new */
1354 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001355};
1356
1357
1358/* islice object ************************************************************/
1359
1360typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001361 PyObject_HEAD
1362 PyObject *it;
1363 Py_ssize_t next;
1364 Py_ssize_t stop;
1365 Py_ssize_t step;
1366 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367} isliceobject;
1368
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001369static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370
1371static PyObject *
1372islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1373{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001374 PyObject *seq;
1375 Py_ssize_t start=0, stop=-1, step=1;
1376 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1377 Py_ssize_t numargs;
1378 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001379
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001380 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1381 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001382
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001383 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1384 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001385
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001386 numargs = PyTuple_Size(args);
1387 if (numargs == 2) {
1388 if (a1 != Py_None) {
1389 stop = PyLong_AsSsize_t(a1);
1390 if (stop == -1) {
1391 if (PyErr_Occurred())
1392 PyErr_Clear();
1393 PyErr_SetString(PyExc_ValueError,
1394 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1395 return NULL;
1396 }
1397 }
1398 } else {
1399 if (a1 != Py_None)
1400 start = PyLong_AsSsize_t(a1);
1401 if (start == -1 && PyErr_Occurred())
1402 PyErr_Clear();
1403 if (a2 != Py_None) {
1404 stop = PyLong_AsSsize_t(a2);
1405 if (stop == -1) {
1406 if (PyErr_Occurred())
1407 PyErr_Clear();
1408 PyErr_SetString(PyExc_ValueError,
1409 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1410 return NULL;
1411 }
1412 }
1413 }
1414 if (start<0 || stop<-1) {
1415 PyErr_SetString(PyExc_ValueError,
1416 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1417 return NULL;
1418 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001419
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001420 if (a3 != NULL) {
1421 if (a3 != Py_None)
1422 step = PyLong_AsSsize_t(a3);
1423 if (step == -1 && PyErr_Occurred())
1424 PyErr_Clear();
1425 }
1426 if (step<1) {
1427 PyErr_SetString(PyExc_ValueError,
1428 "Step for islice() must be a positive integer or None.");
1429 return NULL;
1430 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001431
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001432 /* Get iterator. */
1433 it = PyObject_GetIter(seq);
1434 if (it == NULL)
1435 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001437 /* create isliceobject structure */
1438 lz = (isliceobject *)type->tp_alloc(type, 0);
1439 if (lz == NULL) {
1440 Py_DECREF(it);
1441 return NULL;
1442 }
1443 lz->it = it;
1444 lz->next = start;
1445 lz->stop = stop;
1446 lz->step = step;
1447 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001449 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450}
1451
1452static void
1453islice_dealloc(isliceobject *lz)
1454{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001455 PyObject_GC_UnTrack(lz);
1456 Py_XDECREF(lz->it);
1457 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458}
1459
1460static int
1461islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1462{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 Py_VISIT(lz->it);
1464 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465}
1466
1467static PyObject *
1468islice_next(isliceobject *lz)
1469{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001470 PyObject *item;
1471 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001472 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001473 Py_ssize_t oldnext;
1474 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001476 iternext = *Py_TYPE(it)->tp_iternext;
1477 while (lz->cnt < lz->next) {
1478 item = iternext(it);
1479 if (item == NULL)
1480 return NULL;
1481 Py_DECREF(item);
1482 lz->cnt++;
1483 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001484 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001485 return NULL;
1486 item = iternext(it);
1487 if (item == NULL)
1488 return NULL;
1489 lz->cnt++;
1490 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001491 /* The (size_t) cast below avoids the danger of undefined
1492 behaviour from signed integer overflow. */
1493 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001494 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1495 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001496 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001497}
1498
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001499static PyObject *
1500islice_reduce(isliceobject *lz)
1501{
1502 /* When unpickled, generate a new object with the same bounds,
1503 * then 'setstate' with the next and count
1504 */
1505 PyObject *stop;
1506 if (lz->stop == -1) {
1507 stop = Py_None;
1508 Py_INCREF(stop);
1509 } else {
1510 stop = PyLong_FromSsize_t(lz->stop);
1511 if (stop == NULL)
1512 return NULL;
1513 }
1514 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1515 lz->it, lz->next, stop, lz->step,
1516 lz->cnt);
1517}
1518
1519static PyObject *
1520islice_setstate(isliceobject *lz, PyObject *state)
1521{
1522 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1523 if (cnt == -1 && PyErr_Occurred())
1524 return NULL;
1525 lz->cnt = cnt;
1526 Py_RETURN_NONE;
1527}
1528
1529static PyMethodDef islice_methods[] = {
1530 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1531 reduce_doc},
1532 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1533 setstate_doc},
1534 {NULL, NULL} /* sentinel */
1535};
1536
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537PyDoc_STRVAR(islice_doc,
1538"islice(iterable, [start,] stop [, step]) --> islice object\n\
1539\n\
1540Return an iterator whose next() method returns selected values from an\n\
1541iterable. If start is specified, will skip all preceding elements;\n\
1542otherwise, start defaults to zero. Step defaults to one. If\n\
1543specified as another value, step determines how many values are \n\
1544skipped between successive calls. Works like a slice() on a list\n\
1545but returns an iterator.");
1546
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001547static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001548 PyVarObject_HEAD_INIT(NULL, 0)
1549 "itertools.islice", /* tp_name */
1550 sizeof(isliceobject), /* tp_basicsize */
1551 0, /* tp_itemsize */
1552 /* methods */
1553 (destructor)islice_dealloc, /* tp_dealloc */
1554 0, /* tp_print */
1555 0, /* tp_getattr */
1556 0, /* tp_setattr */
1557 0, /* tp_reserved */
1558 0, /* tp_repr */
1559 0, /* tp_as_number */
1560 0, /* tp_as_sequence */
1561 0, /* tp_as_mapping */
1562 0, /* tp_hash */
1563 0, /* tp_call */
1564 0, /* tp_str */
1565 PyObject_GenericGetAttr, /* tp_getattro */
1566 0, /* tp_setattro */
1567 0, /* tp_as_buffer */
1568 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1569 Py_TPFLAGS_BASETYPE, /* tp_flags */
1570 islice_doc, /* tp_doc */
1571 (traverseproc)islice_traverse, /* tp_traverse */
1572 0, /* tp_clear */
1573 0, /* tp_richcompare */
1574 0, /* tp_weaklistoffset */
1575 PyObject_SelfIter, /* tp_iter */
1576 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001577 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001578 0, /* tp_members */
1579 0, /* tp_getset */
1580 0, /* tp_base */
1581 0, /* tp_dict */
1582 0, /* tp_descr_get */
1583 0, /* tp_descr_set */
1584 0, /* tp_dictoffset */
1585 0, /* tp_init */
1586 0, /* tp_alloc */
1587 islice_new, /* tp_new */
1588 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589};
1590
1591
1592/* starmap object ************************************************************/
1593
1594typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001595 PyObject_HEAD
1596 PyObject *func;
1597 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001598} starmapobject;
1599
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001600static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001601
1602static PyObject *
1603starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1604{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001605 PyObject *func, *seq;
1606 PyObject *it;
1607 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001609 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1610 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001611
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001612 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1613 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001614
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001615 /* Get iterator. */
1616 it = PyObject_GetIter(seq);
1617 if (it == NULL)
1618 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001620 /* create starmapobject structure */
1621 lz = (starmapobject *)type->tp_alloc(type, 0);
1622 if (lz == NULL) {
1623 Py_DECREF(it);
1624 return NULL;
1625 }
1626 Py_INCREF(func);
1627 lz->func = func;
1628 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001630 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001631}
1632
1633static void
1634starmap_dealloc(starmapobject *lz)
1635{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001636 PyObject_GC_UnTrack(lz);
1637 Py_XDECREF(lz->func);
1638 Py_XDECREF(lz->it);
1639 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640}
1641
1642static int
1643starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1644{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001645 Py_VISIT(lz->it);
1646 Py_VISIT(lz->func);
1647 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648}
1649
1650static PyObject *
1651starmap_next(starmapobject *lz)
1652{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001653 PyObject *args;
1654 PyObject *result;
1655 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001656
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001657 args = (*Py_TYPE(it)->tp_iternext)(it);
1658 if (args == NULL)
1659 return NULL;
1660 if (!PyTuple_CheckExact(args)) {
1661 PyObject *newargs = PySequence_Tuple(args);
1662 Py_DECREF(args);
1663 if (newargs == NULL)
1664 return NULL;
1665 args = newargs;
1666 }
1667 result = PyObject_Call(lz->func, args, NULL);
1668 Py_DECREF(args);
1669 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670}
1671
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001672static PyObject *
1673starmap_reduce(starmapobject *lz)
1674{
1675 /* Just pickle the iterator */
1676 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1677}
1678
1679static PyMethodDef starmap_methods[] = {
1680 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1681 reduce_doc},
1682 {NULL, NULL} /* sentinel */
1683};
1684
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001685PyDoc_STRVAR(starmap_doc,
1686"starmap(function, sequence) --> starmap object\n\
1687\n\
1688Return an iterator whose values are returned from the function evaluated\n\
1689with a argument tuple taken from the given sequence.");
1690
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001691static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001692 PyVarObject_HEAD_INIT(NULL, 0)
1693 "itertools.starmap", /* tp_name */
1694 sizeof(starmapobject), /* tp_basicsize */
1695 0, /* tp_itemsize */
1696 /* methods */
1697 (destructor)starmap_dealloc, /* tp_dealloc */
1698 0, /* tp_print */
1699 0, /* tp_getattr */
1700 0, /* tp_setattr */
1701 0, /* tp_reserved */
1702 0, /* tp_repr */
1703 0, /* tp_as_number */
1704 0, /* tp_as_sequence */
1705 0, /* tp_as_mapping */
1706 0, /* tp_hash */
1707 0, /* tp_call */
1708 0, /* tp_str */
1709 PyObject_GenericGetAttr, /* tp_getattro */
1710 0, /* tp_setattro */
1711 0, /* tp_as_buffer */
1712 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1713 Py_TPFLAGS_BASETYPE, /* tp_flags */
1714 starmap_doc, /* tp_doc */
1715 (traverseproc)starmap_traverse, /* tp_traverse */
1716 0, /* tp_clear */
1717 0, /* tp_richcompare */
1718 0, /* tp_weaklistoffset */
1719 PyObject_SelfIter, /* tp_iter */
1720 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001721 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001722 0, /* tp_members */
1723 0, /* tp_getset */
1724 0, /* tp_base */
1725 0, /* tp_dict */
1726 0, /* tp_descr_get */
1727 0, /* tp_descr_set */
1728 0, /* tp_dictoffset */
1729 0, /* tp_init */
1730 0, /* tp_alloc */
1731 starmap_new, /* tp_new */
1732 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001733};
1734
1735
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001736/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737
1738typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001739 PyObject_HEAD
1740 PyObject *source; /* Iterator over input iterables */
1741 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001742} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001743
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001744static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001746static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001747chain_new_internal(PyTypeObject *type, PyObject *source)
1748{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001749 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001751 lz = (chainobject *)type->tp_alloc(type, 0);
1752 if (lz == NULL) {
1753 Py_DECREF(source);
1754 return NULL;
1755 }
1756
1757 lz->source = source;
1758 lz->active = NULL;
1759 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001760}
1761
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001762static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001763chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001765 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001766
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001767 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1768 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001769
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001770 source = PyObject_GetIter(args);
1771 if (source == NULL)
1772 return NULL;
1773
1774 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001775}
1776
1777static PyObject *
1778chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001780 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001781
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001782 source = PyObject_GetIter(arg);
1783 if (source == NULL)
1784 return NULL;
1785
1786 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001787}
1788
1789static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001790chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001791{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001792 PyObject_GC_UnTrack(lz);
1793 Py_XDECREF(lz->active);
1794 Py_XDECREF(lz->source);
1795 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001796}
1797
Raymond Hettinger2012f172003-02-07 05:32:58 +00001798static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001799chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001800{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001801 Py_VISIT(lz->source);
1802 Py_VISIT(lz->active);
1803 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001804}
1805
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001807chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001809 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001810
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001811 if (lz->source == NULL)
1812 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001813
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001814 if (lz->active == NULL) {
1815 PyObject *iterable = PyIter_Next(lz->source);
1816 if (iterable == NULL) {
1817 Py_CLEAR(lz->source);
1818 return NULL; /* no more input sources */
1819 }
1820 lz->active = PyObject_GetIter(iterable);
1821 Py_DECREF(iterable);
1822 if (lz->active == NULL) {
1823 Py_CLEAR(lz->source);
1824 return NULL; /* input not iterable */
1825 }
1826 }
1827 item = PyIter_Next(lz->active);
1828 if (item != NULL)
1829 return item;
1830 if (PyErr_Occurred()) {
1831 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1832 PyErr_Clear();
1833 else
1834 return NULL; /* input raised an exception */
1835 }
1836 Py_CLEAR(lz->active);
1837 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001838}
1839
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001840static PyObject *
1841chain_reduce(chainobject *lz)
1842{
1843 if (lz->source) {
1844 /* we can't pickle function objects (itertools.from_iterable) so
1845 * we must use setstate to replace the iterable. One day we
1846 * will fix pickling of functions
1847 */
1848 if (lz->active) {
1849 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1850 } else {
1851 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1852 }
1853 } else {
1854 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1855 }
1856 return NULL;
1857}
1858
1859static PyObject *
1860chain_setstate(chainobject *lz, PyObject *state)
1861{
1862 PyObject *source, *active=NULL;
1863 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1864 return NULL;
1865
1866 Py_CLEAR(lz->source);
1867 lz->source = source;
1868 Py_INCREF(lz->source);
1869 Py_CLEAR(lz->active);
1870 lz->active = active;
1871 Py_XINCREF(lz->active);
1872 Py_RETURN_NONE;
1873}
1874
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001875PyDoc_STRVAR(chain_doc,
1876"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001877\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001878Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001879first iterable until it is exhausted, then elements from the next\n\
1880iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001881
Christian Heimesf16baeb2008-02-29 14:57:44 +00001882PyDoc_STRVAR(chain_from_iterable_doc,
1883"chain.from_iterable(iterable) --> chain object\n\
1884\n\
1885Alternate chain() contructor taking a single iterable argument\n\
1886that evaluates lazily.");
1887
1888static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001889 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1890 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001891 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1892 reduce_doc},
1893 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1894 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001895 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001896};
1897
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001898static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001899 PyVarObject_HEAD_INIT(NULL, 0)
1900 "itertools.chain", /* tp_name */
1901 sizeof(chainobject), /* tp_basicsize */
1902 0, /* tp_itemsize */
1903 /* methods */
1904 (destructor)chain_dealloc, /* tp_dealloc */
1905 0, /* tp_print */
1906 0, /* tp_getattr */
1907 0, /* tp_setattr */
1908 0, /* tp_reserved */
1909 0, /* tp_repr */
1910 0, /* tp_as_number */
1911 0, /* tp_as_sequence */
1912 0, /* tp_as_mapping */
1913 0, /* tp_hash */
1914 0, /* tp_call */
1915 0, /* tp_str */
1916 PyObject_GenericGetAttr, /* tp_getattro */
1917 0, /* tp_setattro */
1918 0, /* tp_as_buffer */
1919 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1920 Py_TPFLAGS_BASETYPE, /* tp_flags */
1921 chain_doc, /* tp_doc */
1922 (traverseproc)chain_traverse, /* tp_traverse */
1923 0, /* tp_clear */
1924 0, /* tp_richcompare */
1925 0, /* tp_weaklistoffset */
1926 PyObject_SelfIter, /* tp_iter */
1927 (iternextfunc)chain_next, /* tp_iternext */
1928 chain_methods, /* tp_methods */
1929 0, /* tp_members */
1930 0, /* tp_getset */
1931 0, /* tp_base */
1932 0, /* tp_dict */
1933 0, /* tp_descr_get */
1934 0, /* tp_descr_set */
1935 0, /* tp_dictoffset */
1936 0, /* tp_init */
1937 0, /* tp_alloc */
1938 chain_new, /* tp_new */
1939 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001940};
1941
1942
Christian Heimesc3f30c42008-02-22 16:37:40 +00001943/* product object ************************************************************/
1944
1945typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001946 PyObject_HEAD
1947 PyObject *pools; /* tuple of pool tuples */
1948 Py_ssize_t *indices; /* one index per pool */
1949 PyObject *result; /* most recently returned result tuple */
1950 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001951} productobject;
1952
1953static PyTypeObject product_type;
1954
1955static PyObject *
1956product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1957{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001958 productobject *lz;
1959 Py_ssize_t nargs, npools, repeat=1;
1960 PyObject *pools = NULL;
1961 Py_ssize_t *indices = NULL;
1962 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001963
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001964 if (kwds != NULL) {
1965 char *kwlist[] = {"repeat", 0};
1966 PyObject *tmpargs = PyTuple_New(0);
1967 if (tmpargs == NULL)
1968 return NULL;
1969 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1970 Py_DECREF(tmpargs);
1971 return NULL;
1972 }
1973 Py_DECREF(tmpargs);
1974 if (repeat < 0) {
1975 PyErr_SetString(PyExc_ValueError,
1976 "repeat argument cannot be negative");
1977 return NULL;
1978 }
1979 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001980
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001981 assert(PyTuple_Check(args));
1982 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1983 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001984
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001985 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1986 if (indices == NULL) {
1987 PyErr_NoMemory();
1988 goto error;
1989 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001991 pools = PyTuple_New(npools);
1992 if (pools == NULL)
1993 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001994
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001995 for (i=0; i < nargs ; ++i) {
1996 PyObject *item = PyTuple_GET_ITEM(args, i);
1997 PyObject *pool = PySequence_Tuple(item);
1998 if (pool == NULL)
1999 goto error;
2000 PyTuple_SET_ITEM(pools, i, pool);
2001 indices[i] = 0;
2002 }
2003 for ( ; i < npools; ++i) {
2004 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2005 Py_INCREF(pool);
2006 PyTuple_SET_ITEM(pools, i, pool);
2007 indices[i] = 0;
2008 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002010 /* create productobject structure */
2011 lz = (productobject *)type->tp_alloc(type, 0);
2012 if (lz == NULL)
2013 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002014
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002015 lz->pools = pools;
2016 lz->indices = indices;
2017 lz->result = NULL;
2018 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002019
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002020 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002021
2022error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002023 if (indices != NULL)
2024 PyMem_Free(indices);
2025 Py_XDECREF(pools);
2026 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002027}
2028
2029static void
2030product_dealloc(productobject *lz)
2031{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002032 PyObject_GC_UnTrack(lz);
2033 Py_XDECREF(lz->pools);
2034 Py_XDECREF(lz->result);
2035 if (lz->indices != NULL)
2036 PyMem_Free(lz->indices);
2037 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002038}
2039
2040static int
2041product_traverse(productobject *lz, visitproc visit, void *arg)
2042{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002043 Py_VISIT(lz->pools);
2044 Py_VISIT(lz->result);
2045 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002046}
2047
2048static PyObject *
2049product_next(productobject *lz)
2050{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002051 PyObject *pool;
2052 PyObject *elem;
2053 PyObject *oldelem;
2054 PyObject *pools = lz->pools;
2055 PyObject *result = lz->result;
2056 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2057 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002058
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002059 if (lz->stopped)
2060 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002061
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002062 if (result == NULL) {
2063 /* On the first pass, return an initial tuple filled with the
2064 first element from each pool. */
2065 result = PyTuple_New(npools);
2066 if (result == NULL)
2067 goto empty;
2068 lz->result = result;
2069 for (i=0; i < npools; i++) {
2070 pool = PyTuple_GET_ITEM(pools, i);
2071 if (PyTuple_GET_SIZE(pool) == 0)
2072 goto empty;
2073 elem = PyTuple_GET_ITEM(pool, 0);
2074 Py_INCREF(elem);
2075 PyTuple_SET_ITEM(result, i, elem);
2076 }
2077 } else {
2078 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002079
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002080 /* Copy the previous result tuple or re-use it if available */
2081 if (Py_REFCNT(result) > 1) {
2082 PyObject *old_result = result;
2083 result = PyTuple_New(npools);
2084 if (result == NULL)
2085 goto empty;
2086 lz->result = result;
2087 for (i=0; i < npools; i++) {
2088 elem = PyTuple_GET_ITEM(old_result, i);
2089 Py_INCREF(elem);
2090 PyTuple_SET_ITEM(result, i, elem);
2091 }
2092 Py_DECREF(old_result);
2093 }
2094 /* Now, we've got the only copy so we can update it in-place */
2095 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002096
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002097 /* Update the pool indices right-to-left. Only advance to the
2098 next pool when the previous one rolls-over */
2099 for (i=npools-1 ; i >= 0 ; i--) {
2100 pool = PyTuple_GET_ITEM(pools, i);
2101 indices[i]++;
2102 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2103 /* Roll-over and advance to next pool */
2104 indices[i] = 0;
2105 elem = PyTuple_GET_ITEM(pool, 0);
2106 Py_INCREF(elem);
2107 oldelem = PyTuple_GET_ITEM(result, i);
2108 PyTuple_SET_ITEM(result, i, elem);
2109 Py_DECREF(oldelem);
2110 } else {
2111 /* No rollover. Just increment and stop here. */
2112 elem = PyTuple_GET_ITEM(pool, indices[i]);
2113 Py_INCREF(elem);
2114 oldelem = PyTuple_GET_ITEM(result, i);
2115 PyTuple_SET_ITEM(result, i, elem);
2116 Py_DECREF(oldelem);
2117 break;
2118 }
2119 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002121 /* If i is negative, then the indices have all rolled-over
2122 and we're done. */
2123 if (i < 0)
2124 goto empty;
2125 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002127 Py_INCREF(result);
2128 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002129
2130empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002131 lz->stopped = 1;
2132 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002133}
2134
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002135static PyObject *
2136product_reduce(productobject *lz)
2137{
2138 if (lz->stopped) {
2139 return Py_BuildValue("O(())", Py_TYPE(lz));
2140 } else if (lz->result == NULL) {
2141 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2142 } else {
2143 PyObject *indices;
2144 Py_ssize_t n, i;
2145
2146 /* we must pickle the indices use them for setstate, and
2147 * additionally indicate that the iterator has started
2148 */
2149 n = PyTuple_GET_SIZE(lz->pools);
2150 indices = PyTuple_New(n);
2151 if (indices == NULL)
2152 return NULL;
2153 for (i=0; i<n; i++){
2154 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2155 if (!index) {
2156 Py_DECREF(indices);
2157 return NULL;
2158 }
2159 PyTuple_SET_ITEM(indices, i, index);
2160 }
2161 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2162 }
2163}
2164
2165static PyObject *
2166product_setstate(productobject *lz, PyObject *state)
2167{
2168 PyObject *result;
2169 Py_ssize_t n, i;
2170
2171 n = PyTuple_GET_SIZE(lz->pools);
2172 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2173 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2174 return NULL;
2175 }
2176 for (i=0; i<n; i++)
2177 {
2178 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2179 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2180 if (index < 0 && PyErr_Occurred())
2181 return NULL; /* not an integer */
2182 /* clamp the index */
2183 if (index < 0)
2184 index = 0;
2185 else if (index > n-1)
2186 index = n-1;
2187 lz->indices[i] = index;
2188 }
2189
2190 result = PyTuple_New(n);
2191 if (!result)
2192 return NULL;
2193 for (i=0; i<n; i++) {
2194 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2195 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2196 Py_INCREF(element);
2197 PyTuple_SET_ITEM(result, i, element);
2198 }
2199 Py_CLEAR(lz->result);
2200 lz->result = result;
2201 Py_RETURN_NONE;
2202}
2203
2204static PyMethodDef product_methods[] = {
2205 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2206 reduce_doc},
2207 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2208 setstate_doc},
2209 {NULL, NULL} /* sentinel */
2210};
2211
Christian Heimesc3f30c42008-02-22 16:37:40 +00002212PyDoc_STRVAR(product_doc,
2213"product(*iterables) --> product object\n\
2214\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002215Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002216For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2217The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2218cycle in a manner similar to an odometer (with the rightmost element changing\n\
2219on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002220To compute the product of an iterable with itself, specify the number\n\
2221of repetitions with the optional repeat keyword argument. For example,\n\
2222product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002223product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2224product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2225
2226static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002227 PyVarObject_HEAD_INIT(NULL, 0)
2228 "itertools.product", /* tp_name */
2229 sizeof(productobject), /* tp_basicsize */
2230 0, /* tp_itemsize */
2231 /* methods */
2232 (destructor)product_dealloc, /* tp_dealloc */
2233 0, /* tp_print */
2234 0, /* tp_getattr */
2235 0, /* tp_setattr */
2236 0, /* tp_reserved */
2237 0, /* tp_repr */
2238 0, /* tp_as_number */
2239 0, /* tp_as_sequence */
2240 0, /* tp_as_mapping */
2241 0, /* tp_hash */
2242 0, /* tp_call */
2243 0, /* tp_str */
2244 PyObject_GenericGetAttr, /* tp_getattro */
2245 0, /* tp_setattro */
2246 0, /* tp_as_buffer */
2247 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2248 Py_TPFLAGS_BASETYPE, /* tp_flags */
2249 product_doc, /* tp_doc */
2250 (traverseproc)product_traverse, /* tp_traverse */
2251 0, /* tp_clear */
2252 0, /* tp_richcompare */
2253 0, /* tp_weaklistoffset */
2254 PyObject_SelfIter, /* tp_iter */
2255 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002256 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002257 0, /* tp_members */
2258 0, /* tp_getset */
2259 0, /* tp_base */
2260 0, /* tp_dict */
2261 0, /* tp_descr_get */
2262 0, /* tp_descr_set */
2263 0, /* tp_dictoffset */
2264 0, /* tp_init */
2265 0, /* tp_alloc */
2266 product_new, /* tp_new */
2267 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002268};
2269
2270
Christian Heimes380f7f22008-02-28 11:19:05 +00002271/* combinations object ************************************************************/
2272
2273typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002274 PyObject_HEAD
2275 PyObject *pool; /* input converted to a tuple */
2276 Py_ssize_t *indices; /* one index per result element */
2277 PyObject *result; /* most recently returned result tuple */
2278 Py_ssize_t r; /* size of result tuple */
2279 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002280} combinationsobject;
2281
2282static PyTypeObject combinations_type;
2283
2284static PyObject *
2285combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2286{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002287 combinationsobject *co;
2288 Py_ssize_t n;
2289 Py_ssize_t r;
2290 PyObject *pool = NULL;
2291 PyObject *iterable = NULL;
2292 Py_ssize_t *indices = NULL;
2293 Py_ssize_t i;
2294 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002295
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002296 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2297 &iterable, &r))
2298 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002300 pool = PySequence_Tuple(iterable);
2301 if (pool == NULL)
2302 goto error;
2303 n = PyTuple_GET_SIZE(pool);
2304 if (r < 0) {
2305 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2306 goto error;
2307 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002308
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002309 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2310 if (indices == NULL) {
2311 PyErr_NoMemory();
2312 goto error;
2313 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002315 for (i=0 ; i<r ; i++)
2316 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002317
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002318 /* create combinationsobject structure */
2319 co = (combinationsobject *)type->tp_alloc(type, 0);
2320 if (co == NULL)
2321 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002323 co->pool = pool;
2324 co->indices = indices;
2325 co->result = NULL;
2326 co->r = r;
2327 co->stopped = r > n ? 1 : 0;
2328
2329 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002330
2331error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002332 if (indices != NULL)
2333 PyMem_Free(indices);
2334 Py_XDECREF(pool);
2335 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002336}
2337
2338static void
2339combinations_dealloc(combinationsobject *co)
2340{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002341 PyObject_GC_UnTrack(co);
2342 Py_XDECREF(co->pool);
2343 Py_XDECREF(co->result);
2344 if (co->indices != NULL)
2345 PyMem_Free(co->indices);
2346 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002347}
2348
2349static int
2350combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2351{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002352 Py_VISIT(co->pool);
2353 Py_VISIT(co->result);
2354 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002355}
2356
2357static PyObject *
2358combinations_next(combinationsobject *co)
2359{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002360 PyObject *elem;
2361 PyObject *oldelem;
2362 PyObject *pool = co->pool;
2363 Py_ssize_t *indices = co->indices;
2364 PyObject *result = co->result;
2365 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2366 Py_ssize_t r = co->r;
2367 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002368
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002369 if (co->stopped)
2370 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002372 if (result == NULL) {
2373 /* On the first pass, initialize result tuple using the indices */
2374 result = PyTuple_New(r);
2375 if (result == NULL)
2376 goto empty;
2377 co->result = result;
2378 for (i=0; i<r ; i++) {
2379 index = indices[i];
2380 elem = PyTuple_GET_ITEM(pool, index);
2381 Py_INCREF(elem);
2382 PyTuple_SET_ITEM(result, i, elem);
2383 }
2384 } else {
2385 /* Copy the previous result tuple or re-use it if available */
2386 if (Py_REFCNT(result) > 1) {
2387 PyObject *old_result = result;
2388 result = PyTuple_New(r);
2389 if (result == NULL)
2390 goto empty;
2391 co->result = result;
2392 for (i=0; i<r ; i++) {
2393 elem = PyTuple_GET_ITEM(old_result, i);
2394 Py_INCREF(elem);
2395 PyTuple_SET_ITEM(result, i, elem);
2396 }
2397 Py_DECREF(old_result);
2398 }
2399 /* Now, we've got the only copy so we can update it in-place
2400 * CPython's empty tuple is a singleton and cached in
2401 * PyTuple's freelist.
2402 */
2403 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002404
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002405 /* Scan indices right-to-left until finding one that is not
2406 at its maximum (i + n - r). */
2407 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2408 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002409
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002410 /* If i is negative, then the indices are all at
2411 their maximum value and we're done. */
2412 if (i < 0)
2413 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002414
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002415 /* Increment the current index which we know is not at its
2416 maximum. Then move back to the right setting each index
2417 to its lowest possible value (one higher than the index
2418 to its left -- this maintains the sort order invariant). */
2419 indices[i]++;
2420 for (j=i+1 ; j<r ; j++)
2421 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002423 /* Update the result tuple for the new indices
2424 starting with i, the leftmost index that changed */
2425 for ( ; i<r ; i++) {
2426 index = indices[i];
2427 elem = PyTuple_GET_ITEM(pool, index);
2428 Py_INCREF(elem);
2429 oldelem = PyTuple_GET_ITEM(result, i);
2430 PyTuple_SET_ITEM(result, i, elem);
2431 Py_DECREF(oldelem);
2432 }
2433 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002434
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 Py_INCREF(result);
2436 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002437
2438empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002439 co->stopped = 1;
2440 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002441}
2442
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002443static PyObject *
2444combinations_reduce(combinationsobject *lz)
2445{
2446 if (lz->result == NULL) {
2447 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2448 } else if (lz->stopped) {
2449 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2450 } else {
2451 PyObject *indices;
2452 Py_ssize_t i;
2453
2454 /* we must pickle the indices and use them for setstate */
2455 indices = PyTuple_New(lz->r);
2456 if (!indices)
2457 return NULL;
2458 for (i=0; i<lz->r; i++)
2459 {
2460 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2461 if (!index) {
2462 Py_DECREF(indices);
2463 return NULL;
2464 }
2465 PyTuple_SET_ITEM(indices, i, index);
2466 }
2467
2468 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2469 }
2470}
2471
2472static PyObject *
2473combinations_setstate(combinationsobject *lz, PyObject *state)
2474{
2475 PyObject *result;
2476 Py_ssize_t i;
2477 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2478
2479 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2480 {
2481 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2482 return NULL;
2483 }
2484
2485 for (i=0; i<lz->r; i++)
2486 {
2487 Py_ssize_t max;
2488 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2489 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2490 if (index == -1 && PyErr_Occurred())
2491 return NULL; /* not an integer */
2492 max = i + n - lz->r;
2493 /* clamp the index (beware of negative max) */
2494 if (index > max)
2495 index = max;
2496 if (index < 0)
2497 index = 0;
2498 lz->indices[i] = index;
2499 }
2500
2501 result = PyTuple_New(lz->r);
2502 if (result == NULL)
2503 return NULL;
2504 for (i=0; i<lz->r; i++) {
2505 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2506 Py_INCREF(element);
2507 PyTuple_SET_ITEM(result, i, element);
2508 }
2509
2510 Py_CLEAR(lz->result);
2511 lz->result = result;
2512 Py_RETURN_NONE;
2513}
2514
2515static PyMethodDef combinations_methods[] = {
2516 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2517 reduce_doc},
2518 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2519 setstate_doc},
2520 {NULL, NULL} /* sentinel */
2521};
2522
Christian Heimes380f7f22008-02-28 11:19:05 +00002523PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002524"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002525\n\
2526Return successive r-length combinations of elements in the iterable.\n\n\
2527combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2528
2529static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002530 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002531 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002532 sizeof(combinationsobject), /* tp_basicsize */
2533 0, /* tp_itemsize */
2534 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002535 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002536 0, /* tp_print */
2537 0, /* tp_getattr */
2538 0, /* tp_setattr */
2539 0, /* tp_reserved */
2540 0, /* tp_repr */
2541 0, /* tp_as_number */
2542 0, /* tp_as_sequence */
2543 0, /* tp_as_mapping */
2544 0, /* tp_hash */
2545 0, /* tp_call */
2546 0, /* tp_str */
2547 PyObject_GenericGetAttr, /* tp_getattro */
2548 0, /* tp_setattro */
2549 0, /* tp_as_buffer */
2550 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2551 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002552 combinations_doc, /* tp_doc */
2553 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002554 0, /* tp_clear */
2555 0, /* tp_richcompare */
2556 0, /* tp_weaklistoffset */
2557 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002558 (iternextfunc)combinations_next, /* tp_iternext */
2559 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002560 0, /* tp_members */
2561 0, /* tp_getset */
2562 0, /* tp_base */
2563 0, /* tp_dict */
2564 0, /* tp_descr_get */
2565 0, /* tp_descr_set */
2566 0, /* tp_dictoffset */
2567 0, /* tp_init */
2568 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002569 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002570 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002571};
2572
2573
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002574/* combinations with replacement object *******************************************/
2575
2576/* Equivalent to:
2577
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002578 def combinations_with_replacement(iterable, r):
2579 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2580 # number items returned: (n+r-1)! / r! / (n-1)!
2581 pool = tuple(iterable)
2582 n = len(pool)
2583 indices = [0] * r
2584 yield tuple(pool[i] for i in indices)
2585 while 1:
2586 for i in reversed(range(r)):
2587 if indices[i] != n - 1:
2588 break
2589 else:
2590 return
2591 indices[i:] = [indices[i] + 1] * (r - i)
2592 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002593
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002594 def combinations_with_replacement2(iterable, r):
2595 'Alternate version that filters from product()'
2596 pool = tuple(iterable)
2597 n = len(pool)
2598 for indices in product(range(n), repeat=r):
2599 if sorted(indices) == list(indices):
2600 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002601*/
2602typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002603 PyObject_HEAD
2604 PyObject *pool; /* input converted to a tuple */
2605 Py_ssize_t *indices; /* one index per result element */
2606 PyObject *result; /* most recently returned result tuple */
2607 Py_ssize_t r; /* size of result tuple */
2608 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002609} cwrobject;
2610
2611static PyTypeObject cwr_type;
2612
2613static PyObject *
2614cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2615{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002616 cwrobject *co;
2617 Py_ssize_t n;
2618 Py_ssize_t r;
2619 PyObject *pool = NULL;
2620 PyObject *iterable = NULL;
2621 Py_ssize_t *indices = NULL;
2622 Py_ssize_t i;
2623 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002624
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002625 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2626 &iterable, &r))
2627 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002629 pool = PySequence_Tuple(iterable);
2630 if (pool == NULL)
2631 goto error;
2632 n = PyTuple_GET_SIZE(pool);
2633 if (r < 0) {
2634 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2635 goto error;
2636 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002637
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002638 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2639 if (indices == NULL) {
2640 PyErr_NoMemory();
2641 goto error;
2642 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002643
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002644 for (i=0 ; i<r ; i++)
2645 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002646
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002647 /* create cwrobject structure */
2648 co = (cwrobject *)type->tp_alloc(type, 0);
2649 if (co == NULL)
2650 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002651
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002652 co->pool = pool;
2653 co->indices = indices;
2654 co->result = NULL;
2655 co->r = r;
2656 co->stopped = !n && r;
2657
2658 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002659
2660error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002661 if (indices != NULL)
2662 PyMem_Free(indices);
2663 Py_XDECREF(pool);
2664 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002665}
2666
2667static void
2668cwr_dealloc(cwrobject *co)
2669{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002670 PyObject_GC_UnTrack(co);
2671 Py_XDECREF(co->pool);
2672 Py_XDECREF(co->result);
2673 if (co->indices != NULL)
2674 PyMem_Free(co->indices);
2675 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002676}
2677
2678static int
2679cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2680{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002681 Py_VISIT(co->pool);
2682 Py_VISIT(co->result);
2683 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002684}
2685
2686static PyObject *
2687cwr_next(cwrobject *co)
2688{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002689 PyObject *elem;
2690 PyObject *oldelem;
2691 PyObject *pool = co->pool;
2692 Py_ssize_t *indices = co->indices;
2693 PyObject *result = co->result;
2694 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2695 Py_ssize_t r = co->r;
2696 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 if (co->stopped)
2699 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002700
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002701 if (result == NULL) {
2702 /* On the first pass, initialize result tuple using the indices */
2703 result = PyTuple_New(r);
2704 if (result == NULL)
2705 goto empty;
2706 co->result = result;
2707 for (i=0; i<r ; i++) {
2708 index = indices[i];
2709 elem = PyTuple_GET_ITEM(pool, index);
2710 Py_INCREF(elem);
2711 PyTuple_SET_ITEM(result, i, elem);
2712 }
2713 } else {
2714 /* Copy the previous result tuple or re-use it if available */
2715 if (Py_REFCNT(result) > 1) {
2716 PyObject *old_result = result;
2717 result = PyTuple_New(r);
2718 if (result == NULL)
2719 goto empty;
2720 co->result = result;
2721 for (i=0; i<r ; i++) {
2722 elem = PyTuple_GET_ITEM(old_result, i);
2723 Py_INCREF(elem);
2724 PyTuple_SET_ITEM(result, i, elem);
2725 }
2726 Py_DECREF(old_result);
2727 }
2728 /* Now, we've got the only copy so we can update it in-place CPython's
2729 empty tuple is a singleton and cached in PyTuple's freelist. */
2730 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002731
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002732 /* Scan indices right-to-left until finding one that is not
2733 * at its maximum (n-1). */
2734 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2735 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002736
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002737 /* If i is negative, then the indices are all at
2738 their maximum value and we're done. */
2739 if (i < 0)
2740 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002741
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002742 /* Increment the current index which we know is not at its
2743 maximum. Then set all to the right to the same value. */
2744 indices[i]++;
2745 for (j=i+1 ; j<r ; j++)
2746 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002747
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 /* Update the result tuple for the new indices
2749 starting with i, the leftmost index that changed */
2750 for ( ; i<r ; i++) {
2751 index = indices[i];
2752 elem = PyTuple_GET_ITEM(pool, index);
2753 Py_INCREF(elem);
2754 oldelem = PyTuple_GET_ITEM(result, i);
2755 PyTuple_SET_ITEM(result, i, elem);
2756 Py_DECREF(oldelem);
2757 }
2758 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 Py_INCREF(result);
2761 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002762
2763empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002764 co->stopped = 1;
2765 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002766}
2767
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002768static PyObject *
2769cwr_reduce(cwrobject *lz)
2770{
2771 if (lz->result == NULL) {
2772 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2773 } else if (lz->stopped) {
2774 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2775 } else {
2776 PyObject *indices;
2777 Py_ssize_t i;
2778
2779 /* we must pickle the indices and use them for setstate */
2780 indices = PyTuple_New(lz->r);
2781 if (!indices)
2782 return NULL;
2783 for (i=0; i<lz->r; i++)
2784 {
2785 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2786 if (!index) {
2787 Py_DECREF(indices);
2788 return NULL;
2789 }
2790 PyTuple_SET_ITEM(indices, i, index);
2791 }
2792
2793 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2794 }
2795}
2796
2797static PyObject *
2798cwr_setstate(cwrobject *lz, PyObject *state)
2799{
2800 PyObject *result;
2801 Py_ssize_t n, i;
2802
2803 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2804 {
2805 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2806 return NULL;
2807 }
2808
2809 n = PyTuple_GET_SIZE(lz->pool);
2810 for (i=0; i<lz->r; i++)
2811 {
2812 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2813 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2814 if (index < 0 && PyErr_Occurred())
2815 return NULL; /* not an integer */
2816 /* clamp the index */
2817 if (index < 0)
2818 index = 0;
2819 else if (index > n-1)
2820 index = n-1;
2821 lz->indices[i] = index;
2822 }
2823 result = PyTuple_New(lz->r);
2824 if (result == NULL)
2825 return NULL;
2826 for (i=0; i<lz->r; i++) {
2827 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2828 Py_INCREF(element);
2829 PyTuple_SET_ITEM(result, i, element);
2830 }
2831 Py_CLEAR(lz->result);
2832 lz->result = result;
2833 Py_RETURN_NONE;
2834}
2835
2836static PyMethodDef cwr_methods[] = {
2837 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2838 reduce_doc},
2839 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2840 setstate_doc},
2841 {NULL, NULL} /* sentinel */
2842};
2843
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002844PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002845"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002846\n\
2847Return successive r-length combinations of elements in the iterable\n\
2848allowing individual elements to have successive repeats.\n\
2849combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2850
2851static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002852 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002853 "itertools.combinations_with_replacement", /* tp_name */
2854 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002855 0, /* tp_itemsize */
2856 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002857 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002858 0, /* tp_print */
2859 0, /* tp_getattr */
2860 0, /* tp_setattr */
2861 0, /* tp_reserved */
2862 0, /* tp_repr */
2863 0, /* tp_as_number */
2864 0, /* tp_as_sequence */
2865 0, /* tp_as_mapping */
2866 0, /* tp_hash */
2867 0, /* tp_call */
2868 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002869 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002870 0, /* tp_setattro */
2871 0, /* tp_as_buffer */
2872 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002873 Py_TPFLAGS_BASETYPE, /* tp_flags */
2874 cwr_doc, /* tp_doc */
2875 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002876 0, /* tp_clear */
2877 0, /* tp_richcompare */
2878 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002879 PyObject_SelfIter, /* tp_iter */
2880 (iternextfunc)cwr_next, /* tp_iternext */
2881 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002882 0, /* tp_members */
2883 0, /* tp_getset */
2884 0, /* tp_base */
2885 0, /* tp_dict */
2886 0, /* tp_descr_get */
2887 0, /* tp_descr_set */
2888 0, /* tp_dictoffset */
2889 0, /* tp_init */
2890 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002891 cwr_new, /* tp_new */
2892 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002893};
2894
2895
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002896/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002897
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002898def permutations(iterable, r=None):
2899 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2900 pool = tuple(iterable)
2901 n = len(pool)
2902 r = n if r is None else r
2903 indices = range(n)
2904 cycles = range(n-r+1, n+1)[::-1]
2905 yield tuple(pool[i] for i in indices[:r])
2906 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002907 for i in reversed(range(r)):
2908 cycles[i] -= 1
2909 if cycles[i] == 0:
2910 indices[i:] = indices[i+1:] + indices[i:i+1]
2911 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002912 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002913 j = cycles[i]
2914 indices[i], indices[-j] = indices[-j], indices[i]
2915 yield tuple(pool[i] for i in indices[:r])
2916 break
2917 else:
2918 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002919*/
2920
2921typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002922 PyObject_HEAD
2923 PyObject *pool; /* input converted to a tuple */
2924 Py_ssize_t *indices; /* one index per element in the pool */
2925 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2926 PyObject *result; /* most recently returned result tuple */
2927 Py_ssize_t r; /* size of result tuple */
2928 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002929} permutationsobject;
2930
2931static PyTypeObject permutations_type;
2932
2933static PyObject *
2934permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2935{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002936 permutationsobject *po;
2937 Py_ssize_t n;
2938 Py_ssize_t r;
2939 PyObject *robj = Py_None;
2940 PyObject *pool = NULL;
2941 PyObject *iterable = NULL;
2942 Py_ssize_t *indices = NULL;
2943 Py_ssize_t *cycles = NULL;
2944 Py_ssize_t i;
2945 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002947 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2948 &iterable, &robj))
2949 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002950
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002951 pool = PySequence_Tuple(iterable);
2952 if (pool == NULL)
2953 goto error;
2954 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002955
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002956 r = n;
2957 if (robj != Py_None) {
2958 if (!PyLong_Check(robj)) {
2959 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2960 goto error;
2961 }
2962 r = PyLong_AsSsize_t(robj);
2963 if (r == -1 && PyErr_Occurred())
2964 goto error;
2965 }
2966 if (r < 0) {
2967 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2968 goto error;
2969 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002970
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002971 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2972 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2973 if (indices == NULL || cycles == NULL) {
2974 PyErr_NoMemory();
2975 goto error;
2976 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002977
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002978 for (i=0 ; i<n ; i++)
2979 indices[i] = i;
2980 for (i=0 ; i<r ; i++)
2981 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002982
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002983 /* create permutationsobject structure */
2984 po = (permutationsobject *)type->tp_alloc(type, 0);
2985 if (po == NULL)
2986 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002987
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002988 po->pool = pool;
2989 po->indices = indices;
2990 po->cycles = cycles;
2991 po->result = NULL;
2992 po->r = r;
2993 po->stopped = r > n ? 1 : 0;
2994
2995 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002996
2997error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002998 if (indices != NULL)
2999 PyMem_Free(indices);
3000 if (cycles != NULL)
3001 PyMem_Free(cycles);
3002 Py_XDECREF(pool);
3003 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003004}
3005
3006static void
3007permutations_dealloc(permutationsobject *po)
3008{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003009 PyObject_GC_UnTrack(po);
3010 Py_XDECREF(po->pool);
3011 Py_XDECREF(po->result);
3012 PyMem_Free(po->indices);
3013 PyMem_Free(po->cycles);
3014 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003015}
3016
3017static int
3018permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3019{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003020 Py_VISIT(po->pool);
3021 Py_VISIT(po->result);
3022 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003023}
3024
3025static PyObject *
3026permutations_next(permutationsobject *po)
3027{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003028 PyObject *elem;
3029 PyObject *oldelem;
3030 PyObject *pool = po->pool;
3031 Py_ssize_t *indices = po->indices;
3032 Py_ssize_t *cycles = po->cycles;
3033 PyObject *result = po->result;
3034 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3035 Py_ssize_t r = po->r;
3036 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003037
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003038 if (po->stopped)
3039 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003040
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003041 if (result == NULL) {
3042 /* On the first pass, initialize result tuple using the indices */
3043 result = PyTuple_New(r);
3044 if (result == NULL)
3045 goto empty;
3046 po->result = result;
3047 for (i=0; i<r ; i++) {
3048 index = indices[i];
3049 elem = PyTuple_GET_ITEM(pool, index);
3050 Py_INCREF(elem);
3051 PyTuple_SET_ITEM(result, i, elem);
3052 }
3053 } else {
3054 if (n == 0)
3055 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003056
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003057 /* Copy the previous result tuple or re-use it if available */
3058 if (Py_REFCNT(result) > 1) {
3059 PyObject *old_result = result;
3060 result = PyTuple_New(r);
3061 if (result == NULL)
3062 goto empty;
3063 po->result = result;
3064 for (i=0; i<r ; i++) {
3065 elem = PyTuple_GET_ITEM(old_result, i);
3066 Py_INCREF(elem);
3067 PyTuple_SET_ITEM(result, i, elem);
3068 }
3069 Py_DECREF(old_result);
3070 }
3071 /* Now, we've got the only copy so we can update it in-place */
3072 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003073
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003074 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3075 for (i=r-1 ; i>=0 ; i--) {
3076 cycles[i] -= 1;
3077 if (cycles[i] == 0) {
3078 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3079 index = indices[i];
3080 for (j=i ; j<n-1 ; j++)
3081 indices[j] = indices[j+1];
3082 indices[n-1] = index;
3083 cycles[i] = n - i;
3084 } else {
3085 j = cycles[i];
3086 index = indices[i];
3087 indices[i] = indices[n-j];
3088 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003090 for (k=i; k<r ; k++) {
3091 /* start with i, the leftmost element that changed */
3092 /* yield tuple(pool[k] for k in indices[:r]) */
3093 index = indices[k];
3094 elem = PyTuple_GET_ITEM(pool, index);
3095 Py_INCREF(elem);
3096 oldelem = PyTuple_GET_ITEM(result, k);
3097 PyTuple_SET_ITEM(result, k, elem);
3098 Py_DECREF(oldelem);
3099 }
3100 break;
3101 }
3102 }
3103 /* If i is negative, then the cycles have all
3104 rolled-over and we're done. */
3105 if (i < 0)
3106 goto empty;
3107 }
3108 Py_INCREF(result);
3109 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003110
3111empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003112 po->stopped = 1;
3113 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003114}
3115
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003116static PyObject *
3117permutations_reduce(permutationsobject *po)
3118{
3119 if (po->result == NULL) {
3120 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3121 } else if (po->stopped) {
3122 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3123 } else {
3124 PyObject *indices=NULL, *cycles=NULL;
3125 Py_ssize_t n, i;
3126
3127 /* we must pickle the indices and cycles and use them for setstate */
3128 n = PyTuple_GET_SIZE(po->pool);
3129 indices = PyTuple_New(n);
3130 if (indices == NULL)
3131 goto err;
3132 for (i=0; i<n; i++){
3133 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3134 if (!index)
3135 goto err;
3136 PyTuple_SET_ITEM(indices, i, index);
3137 }
3138
3139 cycles = PyTuple_New(po->r);
3140 if (cycles == NULL)
3141 goto err;
3142 for (i=0; i<po->r; i++)
3143 {
3144 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3145 if (!index)
3146 goto err;
3147 PyTuple_SET_ITEM(cycles, i, index);
3148 }
3149 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3150 po->pool, po->r,
3151 indices, cycles);
3152 err:
3153 Py_XDECREF(indices);
3154 Py_XDECREF(cycles);
3155 return NULL;
3156 }
3157}
3158
3159static PyObject *
3160permutations_setstate(permutationsobject *po, PyObject *state)
3161{
3162 PyObject *indices, *cycles, *result;
3163 Py_ssize_t n, i;
3164
3165 if (!PyArg_ParseTuple(state, "O!O!",
3166 &PyTuple_Type, &indices,
3167 &PyTuple_Type, &cycles))
3168 return NULL;
3169
3170 n = PyTuple_GET_SIZE(po->pool);
3171 if (PyTuple_GET_SIZE(indices) != n ||
3172 PyTuple_GET_SIZE(cycles) != po->r)
3173 {
3174 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3175 return NULL;
3176 }
3177
3178 for (i=0; i<n; i++)
3179 {
3180 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3181 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3182 if (index < 0 && PyErr_Occurred())
3183 return NULL; /* not an integer */
3184 /* clamp the index */
3185 if (index < 0)
3186 index = 0;
3187 else if (index > n-1)
3188 index = n-1;
3189 po->indices[i] = index;
3190 }
3191
3192 for (i=0; i<po->r; i++)
3193 {
3194 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3195 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3196 if (index < 0 && PyErr_Occurred())
3197 return NULL; /* not an integer */
3198 if (index < 1)
3199 index = 1;
3200 else if (index > n-i)
3201 index = n-i;
3202 po->cycles[i] = index;
3203 }
3204 result = PyTuple_New(po->r);
3205 if (result == NULL)
3206 return NULL;
3207 for (i=0; i<po->r; i++) {
3208 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3209 Py_INCREF(element);
3210 PyTuple_SET_ITEM(result, i, element);
3211 }
3212 Py_CLEAR(po->result);
3213 po->result = result;
3214 Py_RETURN_NONE;
3215}
3216
3217static PyMethodDef permuations_methods[] = {
3218 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3219 reduce_doc},
3220 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3221 setstate_doc},
3222 {NULL, NULL} /* sentinel */
3223};
3224
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003225PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003226"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003227\n\
3228Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003229permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003230
3231static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003232 PyVarObject_HEAD_INIT(NULL, 0)
3233 "itertools.permutations", /* tp_name */
3234 sizeof(permutationsobject), /* tp_basicsize */
3235 0, /* tp_itemsize */
3236 /* methods */
3237 (destructor)permutations_dealloc, /* tp_dealloc */
3238 0, /* tp_print */
3239 0, /* tp_getattr */
3240 0, /* tp_setattr */
3241 0, /* tp_reserved */
3242 0, /* tp_repr */
3243 0, /* tp_as_number */
3244 0, /* tp_as_sequence */
3245 0, /* tp_as_mapping */
3246 0, /* tp_hash */
3247 0, /* tp_call */
3248 0, /* tp_str */
3249 PyObject_GenericGetAttr, /* tp_getattro */
3250 0, /* tp_setattro */
3251 0, /* tp_as_buffer */
3252 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3253 Py_TPFLAGS_BASETYPE, /* tp_flags */
3254 permutations_doc, /* tp_doc */
3255 (traverseproc)permutations_traverse, /* tp_traverse */
3256 0, /* tp_clear */
3257 0, /* tp_richcompare */
3258 0, /* tp_weaklistoffset */
3259 PyObject_SelfIter, /* tp_iter */
3260 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003261 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003262 0, /* tp_members */
3263 0, /* tp_getset */
3264 0, /* tp_base */
3265 0, /* tp_dict */
3266 0, /* tp_descr_get */
3267 0, /* tp_descr_set */
3268 0, /* tp_dictoffset */
3269 0, /* tp_init */
3270 0, /* tp_alloc */
3271 permutations_new, /* tp_new */
3272 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003273};
3274
Raymond Hettinger482ba772010-12-01 22:48:00 +00003275/* accumulate object ************************************************************/
3276
3277typedef struct {
3278 PyObject_HEAD
3279 PyObject *total;
3280 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003281 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003282} accumulateobject;
3283
3284static PyTypeObject accumulate_type;
3285
3286static PyObject *
3287accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3288{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003289 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003290 PyObject *iterable;
3291 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003292 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003293 accumulateobject *lz;
3294
Raymond Hettinger5d446132011-03-27 18:52:10 -07003295 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3296 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003297 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003298
3299 /* Get iterator. */
3300 it = PyObject_GetIter(iterable);
3301 if (it == NULL)
3302 return NULL;
3303
Raymond Hettinger482ba772010-12-01 22:48:00 +00003304 /* create accumulateobject structure */
3305 lz = (accumulateobject *)type->tp_alloc(type, 0);
3306 if (lz == NULL) {
3307 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003308 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003309 }
3310
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003311 if (binop != Py_None) {
3312 Py_XINCREF(binop);
3313 lz->binop = binop;
3314 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003315 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003316 lz->it = it;
3317 return (PyObject *)lz;
3318}
3319
3320static void
3321accumulate_dealloc(accumulateobject *lz)
3322{
3323 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003324 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003325 Py_XDECREF(lz->total);
3326 Py_XDECREF(lz->it);
3327 Py_TYPE(lz)->tp_free(lz);
3328}
3329
3330static int
3331accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3332{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003333 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003334 Py_VISIT(lz->it);
3335 Py_VISIT(lz->total);
3336 return 0;
3337}
3338
3339static PyObject *
3340accumulate_next(accumulateobject *lz)
3341{
3342 PyObject *val, *oldtotal, *newtotal;
3343
3344 val = PyIter_Next(lz->it);
3345 if (val == NULL)
3346 return NULL;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003347
3348 if (lz->total == NULL) {
3349 Py_INCREF(val);
3350 lz->total = val;
3351 return lz->total;
3352 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003353
3354 if (lz->binop == NULL)
3355 newtotal = PyNumber_Add(lz->total, val);
3356 else
3357 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003358 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003359 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003360 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003361
3362 oldtotal = lz->total;
3363 lz->total = newtotal;
3364 Py_DECREF(oldtotal);
3365
3366 Py_INCREF(newtotal);
3367 return newtotal;
3368}
3369
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003370static PyObject *
3371accumulate_reduce(accumulateobject *lz)
3372{
3373 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3374 lz->it, lz->binop?lz->binop:Py_None,
3375 lz->total?lz->total:Py_None);
3376 }
3377
3378static PyObject *
3379accumulate_setstate(accumulateobject *lz, PyObject *state)
3380{
3381 Py_CLEAR(lz->total);
3382 lz->total = state;
3383 Py_INCREF(lz->total);
3384 Py_RETURN_NONE;
3385}
3386
3387static PyMethodDef accumulate_methods[] = {
3388 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3389 reduce_doc},
3390 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3391 setstate_doc},
3392 {NULL, NULL} /* sentinel */
3393};
3394
Raymond Hettinger482ba772010-12-01 22:48:00 +00003395PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003396"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003397\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003398Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003399
3400static PyTypeObject accumulate_type = {
3401 PyVarObject_HEAD_INIT(NULL, 0)
3402 "itertools.accumulate", /* tp_name */
3403 sizeof(accumulateobject), /* tp_basicsize */
3404 0, /* tp_itemsize */
3405 /* methods */
3406 (destructor)accumulate_dealloc, /* tp_dealloc */
3407 0, /* tp_print */
3408 0, /* tp_getattr */
3409 0, /* tp_setattr */
3410 0, /* tp_reserved */
3411 0, /* tp_repr */
3412 0, /* tp_as_number */
3413 0, /* tp_as_sequence */
3414 0, /* tp_as_mapping */
3415 0, /* tp_hash */
3416 0, /* tp_call */
3417 0, /* tp_str */
3418 PyObject_GenericGetAttr, /* tp_getattro */
3419 0, /* tp_setattro */
3420 0, /* tp_as_buffer */
3421 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3422 Py_TPFLAGS_BASETYPE, /* tp_flags */
3423 accumulate_doc, /* tp_doc */
3424 (traverseproc)accumulate_traverse, /* tp_traverse */
3425 0, /* tp_clear */
3426 0, /* tp_richcompare */
3427 0, /* tp_weaklistoffset */
3428 PyObject_SelfIter, /* tp_iter */
3429 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003430 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003431 0, /* tp_members */
3432 0, /* tp_getset */
3433 0, /* tp_base */
3434 0, /* tp_dict */
3435 0, /* tp_descr_get */
3436 0, /* tp_descr_set */
3437 0, /* tp_dictoffset */
3438 0, /* tp_init */
3439 0, /* tp_alloc */
3440 accumulate_new, /* tp_new */
3441 PyObject_GC_Del, /* tp_free */
3442};
3443
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003444
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003445/* compress object ************************************************************/
3446
3447/* Equivalent to:
3448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003449 def compress(data, selectors):
3450 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3451 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003452*/
3453
3454typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003455 PyObject_HEAD
3456 PyObject *data;
3457 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003458} compressobject;
3459
3460static PyTypeObject compress_type;
3461
3462static PyObject *
3463compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3464{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003465 PyObject *seq1, *seq2;
3466 PyObject *data=NULL, *selectors=NULL;
3467 compressobject *lz;
3468 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003469
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003470 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3471 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003472
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003473 data = PyObject_GetIter(seq1);
3474 if (data == NULL)
3475 goto fail;
3476 selectors = PyObject_GetIter(seq2);
3477 if (selectors == NULL)
3478 goto fail;
3479
3480 /* create compressobject structure */
3481 lz = (compressobject *)type->tp_alloc(type, 0);
3482 if (lz == NULL)
3483 goto fail;
3484 lz->data = data;
3485 lz->selectors = selectors;
3486 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003487
3488fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003489 Py_XDECREF(data);
3490 Py_XDECREF(selectors);
3491 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003492}
3493
3494static void
3495compress_dealloc(compressobject *lz)
3496{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003497 PyObject_GC_UnTrack(lz);
3498 Py_XDECREF(lz->data);
3499 Py_XDECREF(lz->selectors);
3500 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003501}
3502
3503static int
3504compress_traverse(compressobject *lz, visitproc visit, void *arg)
3505{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003506 Py_VISIT(lz->data);
3507 Py_VISIT(lz->selectors);
3508 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003509}
3510
3511static PyObject *
3512compress_next(compressobject *lz)
3513{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003514 PyObject *data = lz->data, *selectors = lz->selectors;
3515 PyObject *datum, *selector;
3516 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3517 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3518 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003519
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003520 while (1) {
3521 /* Steps: get datum, get selector, evaluate selector.
3522 Order is important (to match the pure python version
3523 in terms of which input gets a chance to raise an
3524 exception first).
3525 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003526
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003527 datum = datanext(data);
3528 if (datum == NULL)
3529 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003530
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003531 selector = selectornext(selectors);
3532 if (selector == NULL) {
3533 Py_DECREF(datum);
3534 return NULL;
3535 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003536
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003537 ok = PyObject_IsTrue(selector);
3538 Py_DECREF(selector);
3539 if (ok == 1)
3540 return datum;
3541 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003542 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003543 return NULL;
3544 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003545}
3546
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003547static PyObject *
3548compress_reduce(compressobject *lz)
3549{
3550 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3551 lz->data, lz->selectors);
3552 }
3553
3554static PyMethodDef compress_methods[] = {
3555 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3556 reduce_doc},
3557 {NULL, NULL} /* sentinel */
3558};
3559
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003560PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003561"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003562\n\
3563Return data elements corresponding to true selector elements.\n\
3564Forms a shorter iterator from selected data elements using the\n\
3565selectors to choose the data elements.");
3566
3567static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003568 PyVarObject_HEAD_INIT(NULL, 0)
3569 "itertools.compress", /* tp_name */
3570 sizeof(compressobject), /* tp_basicsize */
3571 0, /* tp_itemsize */
3572 /* methods */
3573 (destructor)compress_dealloc, /* tp_dealloc */
3574 0, /* tp_print */
3575 0, /* tp_getattr */
3576 0, /* tp_setattr */
3577 0, /* tp_reserved */
3578 0, /* tp_repr */
3579 0, /* tp_as_number */
3580 0, /* tp_as_sequence */
3581 0, /* tp_as_mapping */
3582 0, /* tp_hash */
3583 0, /* tp_call */
3584 0, /* tp_str */
3585 PyObject_GenericGetAttr, /* tp_getattro */
3586 0, /* tp_setattro */
3587 0, /* tp_as_buffer */
3588 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3589 Py_TPFLAGS_BASETYPE, /* tp_flags */
3590 compress_doc, /* tp_doc */
3591 (traverseproc)compress_traverse, /* tp_traverse */
3592 0, /* tp_clear */
3593 0, /* tp_richcompare */
3594 0, /* tp_weaklistoffset */
3595 PyObject_SelfIter, /* tp_iter */
3596 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003597 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003598 0, /* tp_members */
3599 0, /* tp_getset */
3600 0, /* tp_base */
3601 0, /* tp_dict */
3602 0, /* tp_descr_get */
3603 0, /* tp_descr_set */
3604 0, /* tp_dictoffset */
3605 0, /* tp_init */
3606 0, /* tp_alloc */
3607 compress_new, /* tp_new */
3608 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003609};
3610
3611
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003612/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003613
3614typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003615 PyObject_HEAD
3616 PyObject *func;
3617 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003618} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003619
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003620static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003621
3622static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003623filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003624{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003625 PyObject *func, *seq;
3626 PyObject *it;
3627 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003629 if (type == &filterfalse_type &&
3630 !_PyArg_NoKeywords("filterfalse()", kwds))
3631 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003632
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003633 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3634 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003635
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003636 /* Get iterator. */
3637 it = PyObject_GetIter(seq);
3638 if (it == NULL)
3639 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003640
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003641 /* create filterfalseobject structure */
3642 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3643 if (lz == NULL) {
3644 Py_DECREF(it);
3645 return NULL;
3646 }
3647 Py_INCREF(func);
3648 lz->func = func;
3649 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003650
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003651 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003652}
3653
3654static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003655filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003656{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003657 PyObject_GC_UnTrack(lz);
3658 Py_XDECREF(lz->func);
3659 Py_XDECREF(lz->it);
3660 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003661}
3662
3663static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003664filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003665{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003666 Py_VISIT(lz->it);
3667 Py_VISIT(lz->func);
3668 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003669}
3670
3671static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003672filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003673{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003674 PyObject *item;
3675 PyObject *it = lz->it;
3676 long ok;
3677 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003679 iternext = *Py_TYPE(it)->tp_iternext;
3680 for (;;) {
3681 item = iternext(it);
3682 if (item == NULL)
3683 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003684
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003685 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3686 ok = PyObject_IsTrue(item);
3687 } else {
3688 PyObject *good;
3689 good = PyObject_CallFunctionObjArgs(lz->func,
3690 item, NULL);
3691 if (good == NULL) {
3692 Py_DECREF(item);
3693 return NULL;
3694 }
3695 ok = PyObject_IsTrue(good);
3696 Py_DECREF(good);
3697 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003698 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003699 return item;
3700 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003701 if (ok < 0)
3702 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003703 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003704}
3705
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003706static PyObject *
3707filterfalse_reduce(filterfalseobject *lz)
3708{
3709 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3710 lz->func, lz->it);
3711 }
3712
3713static PyMethodDef filterfalse_methods[] = {
3714 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3715 reduce_doc},
3716 {NULL, NULL} /* sentinel */
3717};
3718
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003719PyDoc_STRVAR(filterfalse_doc,
3720"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003721\n\
3722Return those items of sequence for which function(item) is false.\n\
3723If function is None, return the items that are false.");
3724
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003725static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003726 PyVarObject_HEAD_INIT(NULL, 0)
3727 "itertools.filterfalse", /* tp_name */
3728 sizeof(filterfalseobject), /* tp_basicsize */
3729 0, /* tp_itemsize */
3730 /* methods */
3731 (destructor)filterfalse_dealloc, /* tp_dealloc */
3732 0, /* tp_print */
3733 0, /* tp_getattr */
3734 0, /* tp_setattr */
3735 0, /* tp_reserved */
3736 0, /* tp_repr */
3737 0, /* tp_as_number */
3738 0, /* tp_as_sequence */
3739 0, /* tp_as_mapping */
3740 0, /* tp_hash */
3741 0, /* tp_call */
3742 0, /* tp_str */
3743 PyObject_GenericGetAttr, /* tp_getattro */
3744 0, /* tp_setattro */
3745 0, /* tp_as_buffer */
3746 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3747 Py_TPFLAGS_BASETYPE, /* tp_flags */
3748 filterfalse_doc, /* tp_doc */
3749 (traverseproc)filterfalse_traverse, /* tp_traverse */
3750 0, /* tp_clear */
3751 0, /* tp_richcompare */
3752 0, /* tp_weaklistoffset */
3753 PyObject_SelfIter, /* tp_iter */
3754 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003755 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003756 0, /* tp_members */
3757 0, /* tp_getset */
3758 0, /* tp_base */
3759 0, /* tp_dict */
3760 0, /* tp_descr_get */
3761 0, /* tp_descr_set */
3762 0, /* tp_dictoffset */
3763 0, /* tp_init */
3764 0, /* tp_alloc */
3765 filterfalse_new, /* tp_new */
3766 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003767};
3768
3769
3770/* count object ************************************************************/
3771
3772typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003773 PyObject_HEAD
3774 Py_ssize_t cnt;
3775 PyObject *long_cnt;
3776 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003777} countobject;
3778
Raymond Hettinger30729212009-02-12 06:28:27 +00003779/* Counting logic and invariants:
3780
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003781fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003782
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003783 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3784 Advances with: cnt += 1
3785 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003786
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003787slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003788
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003789 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3790 All counting is done with python objects (no overflows or underflows).
3791 Advances with: long_cnt += long_step
3792 Step may be zero -- effectively a slow version of repeat(cnt).
3793 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003794*/
3795
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003796static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003797
3798static PyObject *
3799count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3800{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003801 countobject *lz;
3802 int slow_mode = 0;
3803 Py_ssize_t cnt = 0;
3804 PyObject *long_cnt = NULL;
3805 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003806 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003807 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003808
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003809 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3810 kwlist, &long_cnt, &long_step))
3811 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003813 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3814 (long_step != NULL && !PyNumber_Check(long_step))) {
3815 PyErr_SetString(PyExc_TypeError, "a number is required");
3816 return NULL;
3817 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003818
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003819 if (long_cnt != NULL) {
3820 cnt = PyLong_AsSsize_t(long_cnt);
3821 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3822 PyErr_Clear();
3823 slow_mode = 1;
3824 }
3825 Py_INCREF(long_cnt);
3826 } else {
3827 cnt = 0;
3828 long_cnt = PyLong_FromLong(0);
3829 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003831 /* If not specified, step defaults to 1 */
3832 if (long_step == NULL) {
3833 long_step = PyLong_FromLong(1);
3834 if (long_step == NULL) {
3835 Py_DECREF(long_cnt);
3836 return NULL;
3837 }
3838 } else
3839 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003840
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003841 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003842
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003843 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003844 step = PyLong_AsLong(long_step);
3845 if (step != 1) {
3846 slow_mode = 1;
3847 if (step == -1 && PyErr_Occurred())
3848 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003849 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003850
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003851 if (slow_mode)
3852 cnt = PY_SSIZE_T_MAX;
3853 else
3854 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003855
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003856 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3857 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3858 assert(slow_mode ||
3859 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003860
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003861 /* create countobject structure */
3862 lz = (countobject *)type->tp_alloc(type, 0);
3863 if (lz == NULL) {
3864 Py_XDECREF(long_cnt);
3865 return NULL;
3866 }
3867 lz->cnt = cnt;
3868 lz->long_cnt = long_cnt;
3869 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003870
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003871 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003872}
3873
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003874static void
3875count_dealloc(countobject *lz)
3876{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003877 PyObject_GC_UnTrack(lz);
3878 Py_XDECREF(lz->long_cnt);
3879 Py_XDECREF(lz->long_step);
3880 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003881}
3882
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003883static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003884count_traverse(countobject *lz, visitproc visit, void *arg)
3885{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003886 Py_VISIT(lz->long_cnt);
3887 Py_VISIT(lz->long_step);
3888 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003889}
3890
3891static PyObject *
3892count_nextlong(countobject *lz)
3893{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003894 PyObject *long_cnt;
3895 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003896
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003897 long_cnt = lz->long_cnt;
3898 if (long_cnt == NULL) {
3899 /* Switch to slow_mode */
3900 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3901 if (long_cnt == NULL)
3902 return NULL;
3903 }
3904 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003905
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003906 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3907 if (stepped_up == NULL)
3908 return NULL;
3909 lz->long_cnt = stepped_up;
3910 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003911}
3912
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003913static PyObject *
3914count_next(countobject *lz)
3915{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003916 if (lz->cnt == PY_SSIZE_T_MAX)
3917 return count_nextlong(lz);
3918 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003919}
3920
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003921static PyObject *
3922count_repr(countobject *lz)
3923{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003924 if (lz->cnt != PY_SSIZE_T_MAX)
3925 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003927 if (PyLong_Check(lz->long_step)) {
3928 long step = PyLong_AsLong(lz->long_step);
3929 if (step == -1 && PyErr_Occurred()) {
3930 PyErr_Clear();
3931 }
3932 if (step == 1) {
3933 /* Don't display step when it is an integer equal to 1 */
3934 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3935 }
3936 }
3937 return PyUnicode_FromFormat("count(%R, %R)",
3938 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003939}
3940
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003941static PyObject *
3942count_reduce(countobject *lz)
3943{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003944 if (lz->cnt == PY_SSIZE_T_MAX)
3945 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3946 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003947}
3948
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003949static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003950 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003951 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003952 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003953};
3954
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003955PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003956 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003957\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003958Return a count object whose .__next__() method returns consecutive values.\n\
3959Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003960 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003961 x = firstval\n\
3962 while 1:\n\
3963 yield x\n\
3964 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003965
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003966static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003967 PyVarObject_HEAD_INIT(NULL, 0)
3968 "itertools.count", /* tp_name */
3969 sizeof(countobject), /* tp_basicsize */
3970 0, /* tp_itemsize */
3971 /* methods */
3972 (destructor)count_dealloc, /* tp_dealloc */
3973 0, /* tp_print */
3974 0, /* tp_getattr */
3975 0, /* tp_setattr */
3976 0, /* tp_reserved */
3977 (reprfunc)count_repr, /* tp_repr */
3978 0, /* tp_as_number */
3979 0, /* tp_as_sequence */
3980 0, /* tp_as_mapping */
3981 0, /* tp_hash */
3982 0, /* tp_call */
3983 0, /* tp_str */
3984 PyObject_GenericGetAttr, /* tp_getattro */
3985 0, /* tp_setattro */
3986 0, /* tp_as_buffer */
3987 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3988 Py_TPFLAGS_BASETYPE, /* tp_flags */
3989 count_doc, /* tp_doc */
3990 (traverseproc)count_traverse, /* tp_traverse */
3991 0, /* tp_clear */
3992 0, /* tp_richcompare */
3993 0, /* tp_weaklistoffset */
3994 PyObject_SelfIter, /* tp_iter */
3995 (iternextfunc)count_next, /* tp_iternext */
3996 count_methods, /* tp_methods */
3997 0, /* tp_members */
3998 0, /* tp_getset */
3999 0, /* tp_base */
4000 0, /* tp_dict */
4001 0, /* tp_descr_get */
4002 0, /* tp_descr_set */
4003 0, /* tp_dictoffset */
4004 0, /* tp_init */
4005 0, /* tp_alloc */
4006 count_new, /* tp_new */
4007 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004008};
4009
4010
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004011/* repeat object ************************************************************/
4012
4013typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004014 PyObject_HEAD
4015 PyObject *element;
4016 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017} repeatobject;
4018
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004019static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004020
4021static PyObject *
4022repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4023{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004024 repeatobject *ro;
4025 PyObject *element;
4026 Py_ssize_t cnt = -1;
4027 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004028
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004029 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4030 &element, &cnt))
4031 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004033 if (PyTuple_Size(args) == 2 && cnt < 0)
4034 cnt = 0;
4035
4036 ro = (repeatobject *)type->tp_alloc(type, 0);
4037 if (ro == NULL)
4038 return NULL;
4039 Py_INCREF(element);
4040 ro->element = element;
4041 ro->cnt = cnt;
4042 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004043}
4044
4045static void
4046repeat_dealloc(repeatobject *ro)
4047{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004048 PyObject_GC_UnTrack(ro);
4049 Py_XDECREF(ro->element);
4050 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004051}
4052
4053static int
4054repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4055{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004056 Py_VISIT(ro->element);
4057 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004058}
4059
4060static PyObject *
4061repeat_next(repeatobject *ro)
4062{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004063 if (ro->cnt == 0)
4064 return NULL;
4065 if (ro->cnt > 0)
4066 ro->cnt--;
4067 Py_INCREF(ro->element);
4068 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004069}
4070
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004071static PyObject *
4072repeat_repr(repeatobject *ro)
4073{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004074 if (ro->cnt == -1)
4075 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4076 else
4077 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4078}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004079
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004080static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004081repeat_len(repeatobject *ro)
4082{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004083 if (ro->cnt == -1) {
4084 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4085 return NULL;
4086 }
4087 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004088}
4089
Armin Rigof5b3e362006-02-11 21:32:43 +00004090PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004091
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004092static PyObject *
4093repeat_reduce(repeatobject *ro)
4094{
4095 /* unpickle this so that a new repeat iterator is constructed with an
4096 * object, then call __setstate__ on it to set cnt
4097 */
4098 if (ro->cnt >= 0)
4099 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4100 else
4101 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4102}
4103
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004104static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004105 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004106 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004107 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004108};
4109
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004110PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004111"repeat(object [,times]) -> create an iterator which returns the object\n\
4112for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004113endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004115static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004116 PyVarObject_HEAD_INIT(NULL, 0)
4117 "itertools.repeat", /* tp_name */
4118 sizeof(repeatobject), /* tp_basicsize */
4119 0, /* tp_itemsize */
4120 /* methods */
4121 (destructor)repeat_dealloc, /* tp_dealloc */
4122 0, /* tp_print */
4123 0, /* tp_getattr */
4124 0, /* tp_setattr */
4125 0, /* tp_reserved */
4126 (reprfunc)repeat_repr, /* tp_repr */
4127 0, /* tp_as_number */
4128 0, /* tp_as_sequence */
4129 0, /* tp_as_mapping */
4130 0, /* tp_hash */
4131 0, /* tp_call */
4132 0, /* tp_str */
4133 PyObject_GenericGetAttr, /* tp_getattro */
4134 0, /* tp_setattro */
4135 0, /* tp_as_buffer */
4136 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4137 Py_TPFLAGS_BASETYPE, /* tp_flags */
4138 repeat_doc, /* tp_doc */
4139 (traverseproc)repeat_traverse, /* tp_traverse */
4140 0, /* tp_clear */
4141 0, /* tp_richcompare */
4142 0, /* tp_weaklistoffset */
4143 PyObject_SelfIter, /* tp_iter */
4144 (iternextfunc)repeat_next, /* tp_iternext */
4145 repeat_methods, /* tp_methods */
4146 0, /* tp_members */
4147 0, /* tp_getset */
4148 0, /* tp_base */
4149 0, /* tp_dict */
4150 0, /* tp_descr_get */
4151 0, /* tp_descr_set */
4152 0, /* tp_dictoffset */
4153 0, /* tp_init */
4154 0, /* tp_alloc */
4155 repeat_new, /* tp_new */
4156 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004157};
4158
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004159/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004160
4161#include "Python.h"
4162
4163typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004164 PyObject_HEAD
4165 Py_ssize_t tuplesize;
4166 Py_ssize_t numactive;
4167 PyObject *ittuple; /* tuple of iterators */
4168 PyObject *result;
4169 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004170} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004171
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004172static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004173
4174static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004175zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004176{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004177 ziplongestobject *lz;
4178 Py_ssize_t i;
4179 PyObject *ittuple; /* tuple of iterators */
4180 PyObject *result;
4181 PyObject *fillvalue = Py_None;
4182 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004183
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004184 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4185 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4186 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4187 PyErr_SetString(PyExc_TypeError,
4188 "zip_longest() got an unexpected keyword argument");
4189 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004190 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004191 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004192
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004193 /* args must be a tuple */
4194 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004195
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 /* obtain iterators */
4197 ittuple = PyTuple_New(tuplesize);
4198 if (ittuple == NULL)
4199 return NULL;
4200 for (i=0; i < tuplesize; ++i) {
4201 PyObject *item = PyTuple_GET_ITEM(args, i);
4202 PyObject *it = PyObject_GetIter(item);
4203 if (it == NULL) {
4204 if (PyErr_ExceptionMatches(PyExc_TypeError))
4205 PyErr_Format(PyExc_TypeError,
4206 "zip_longest argument #%zd must support iteration",
4207 i+1);
4208 Py_DECREF(ittuple);
4209 return NULL;
4210 }
4211 PyTuple_SET_ITEM(ittuple, i, it);
4212 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004213
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004214 /* create a result holder */
4215 result = PyTuple_New(tuplesize);
4216 if (result == NULL) {
4217 Py_DECREF(ittuple);
4218 return NULL;
4219 }
4220 for (i=0 ; i < tuplesize ; i++) {
4221 Py_INCREF(Py_None);
4222 PyTuple_SET_ITEM(result, i, Py_None);
4223 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004224
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004225 /* create ziplongestobject structure */
4226 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4227 if (lz == NULL) {
4228 Py_DECREF(ittuple);
4229 Py_DECREF(result);
4230 return NULL;
4231 }
4232 lz->ittuple = ittuple;
4233 lz->tuplesize = tuplesize;
4234 lz->numactive = tuplesize;
4235 lz->result = result;
4236 Py_INCREF(fillvalue);
4237 lz->fillvalue = fillvalue;
4238 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004239}
4240
4241static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004242zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004243{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004244 PyObject_GC_UnTrack(lz);
4245 Py_XDECREF(lz->ittuple);
4246 Py_XDECREF(lz->result);
4247 Py_XDECREF(lz->fillvalue);
4248 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004249}
4250
4251static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004252zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004253{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004254 Py_VISIT(lz->ittuple);
4255 Py_VISIT(lz->result);
4256 Py_VISIT(lz->fillvalue);
4257 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004258}
4259
4260static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004261zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004262{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004263 Py_ssize_t i;
4264 Py_ssize_t tuplesize = lz->tuplesize;
4265 PyObject *result = lz->result;
4266 PyObject *it;
4267 PyObject *item;
4268 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004269
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004270 if (tuplesize == 0)
4271 return NULL;
4272 if (lz->numactive == 0)
4273 return NULL;
4274 if (Py_REFCNT(result) == 1) {
4275 Py_INCREF(result);
4276 for (i=0 ; i < tuplesize ; i++) {
4277 it = PyTuple_GET_ITEM(lz->ittuple, i);
4278 if (it == NULL) {
4279 Py_INCREF(lz->fillvalue);
4280 item = lz->fillvalue;
4281 } else {
4282 item = PyIter_Next(it);
4283 if (item == NULL) {
4284 lz->numactive -= 1;
4285 if (lz->numactive == 0 || PyErr_Occurred()) {
4286 lz->numactive = 0;
4287 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004288 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004289 } else {
4290 Py_INCREF(lz->fillvalue);
4291 item = lz->fillvalue;
4292 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4293 Py_DECREF(it);
4294 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004295 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004296 }
4297 olditem = PyTuple_GET_ITEM(result, i);
4298 PyTuple_SET_ITEM(result, i, item);
4299 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004300 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004301 } else {
4302 result = PyTuple_New(tuplesize);
4303 if (result == NULL)
4304 return NULL;
4305 for (i=0 ; i < tuplesize ; i++) {
4306 it = PyTuple_GET_ITEM(lz->ittuple, i);
4307 if (it == NULL) {
4308 Py_INCREF(lz->fillvalue);
4309 item = lz->fillvalue;
4310 } else {
4311 item = PyIter_Next(it);
4312 if (item == NULL) {
4313 lz->numactive -= 1;
4314 if (lz->numactive == 0 || PyErr_Occurred()) {
4315 lz->numactive = 0;
4316 Py_DECREF(result);
4317 return NULL;
4318 } else {
4319 Py_INCREF(lz->fillvalue);
4320 item = lz->fillvalue;
4321 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4322 Py_DECREF(it);
4323 }
4324 }
4325 }
4326 PyTuple_SET_ITEM(result, i, item);
4327 }
4328 }
4329 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004330}
4331
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004332static PyObject *
4333zip_longest_reduce(ziplongestobject *lz)
4334{
4335
4336 /* Create a new tuple with empty sequences where appropriate to pickle.
4337 * Then use setstate to set the fillvalue
4338 */
4339 int i;
4340 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4341 if (args == NULL)
4342 return NULL;
4343 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4344 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4345 if (elem == NULL) {
4346 elem = PyTuple_New(0);
4347 if (elem == NULL) {
4348 Py_DECREF(args);
4349 return NULL;
4350 }
4351 } else
4352 Py_INCREF(elem);
4353 PyTuple_SET_ITEM(args, i, elem);
4354 }
4355 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4356}
4357
4358static PyObject *
4359zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4360{
4361 Py_CLEAR(lz->fillvalue);
4362 lz->fillvalue = state;
4363 Py_INCREF(lz->fillvalue);
4364 Py_RETURN_NONE;
4365}
4366
4367static PyMethodDef zip_longest_methods[] = {
4368 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4369 reduce_doc},
4370 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4371 setstate_doc},
4372 {NULL, NULL} /* sentinel */
4373};
4374
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004375PyDoc_STRVAR(zip_longest_doc,
4376"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004377\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004378Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004379the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004380method continues until the longest iterable in the argument sequence\n\
4381is exhausted and then it raises StopIteration. When the shorter iterables\n\
4382are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4383defaults to None or can be specified by a keyword argument.\n\
4384");
4385
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004386static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004387 PyVarObject_HEAD_INIT(NULL, 0)
4388 "itertools.zip_longest", /* tp_name */
4389 sizeof(ziplongestobject), /* tp_basicsize */
4390 0, /* tp_itemsize */
4391 /* methods */
4392 (destructor)zip_longest_dealloc, /* tp_dealloc */
4393 0, /* tp_print */
4394 0, /* tp_getattr */
4395 0, /* tp_setattr */
4396 0, /* tp_reserved */
4397 0, /* tp_repr */
4398 0, /* tp_as_number */
4399 0, /* tp_as_sequence */
4400 0, /* tp_as_mapping */
4401 0, /* tp_hash */
4402 0, /* tp_call */
4403 0, /* tp_str */
4404 PyObject_GenericGetAttr, /* tp_getattro */
4405 0, /* tp_setattro */
4406 0, /* tp_as_buffer */
4407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4408 Py_TPFLAGS_BASETYPE, /* tp_flags */
4409 zip_longest_doc, /* tp_doc */
4410 (traverseproc)zip_longest_traverse, /* tp_traverse */
4411 0, /* tp_clear */
4412 0, /* tp_richcompare */
4413 0, /* tp_weaklistoffset */
4414 PyObject_SelfIter, /* tp_iter */
4415 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004416 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004417 0, /* tp_members */
4418 0, /* tp_getset */
4419 0, /* tp_base */
4420 0, /* tp_dict */
4421 0, /* tp_descr_get */
4422 0, /* tp_descr_set */
4423 0, /* tp_dictoffset */
4424 0, /* tp_init */
4425 0, /* tp_alloc */
4426 zip_longest_new, /* tp_new */
4427 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004428};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004429
4430/* module level code ********************************************************/
4431
4432PyDoc_STRVAR(module_doc,
4433"Functional tools for creating and using iterators.\n\
4434\n\
4435Infinite iterators:\n\
4436count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004437cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004438repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004439\n\
4440Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004441accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004442chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
4443compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4444dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4445groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004446filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004447islice(seq, [start,] stop [, step]) --> elements from\n\
4448 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004449starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004450tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004451takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004452zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004453\n\
4454Combinatoric generators:\n\
4455product(p, q, ... [repeat=1]) --> cartesian product\n\
4456permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004457combinations(p, r)\n\
4458combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004459");
4460
4461
Raymond Hettingerad983e72003-11-12 14:32:26 +00004462static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004463 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4464 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004465};
4466
Martin v. Löwis1a214512008-06-11 05:26:20 +00004467
4468static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004469 PyModuleDef_HEAD_INIT,
4470 "itertools",
4471 module_doc,
4472 -1,
4473 module_methods,
4474 NULL,
4475 NULL,
4476 NULL,
4477 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004478};
4479
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004480PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004481PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004482{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004483 int i;
4484 PyObject *m;
4485 char *name;
4486 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004487 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004488 &combinations_type,
4489 &cwr_type,
4490 &cycle_type,
4491 &dropwhile_type,
4492 &takewhile_type,
4493 &islice_type,
4494 &starmap_type,
4495 &chain_type,
4496 &compress_type,
4497 &filterfalse_type,
4498 &count_type,
4499 &ziplongest_type,
4500 &permutations_type,
4501 &product_type,
4502 &repeat_type,
4503 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004504 &_grouper_type,
4505 &tee_type,
4506 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004507 NULL
4508 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004509
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004510 Py_TYPE(&teedataobject_type) = &PyType_Type;
4511 m = PyModule_Create(&itertoolsmodule);
4512 if (m == NULL)
4513 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004514
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004515 for (i=0 ; typelist[i] != NULL ; i++) {
4516 if (PyType_Ready(typelist[i]) < 0)
4517 return NULL;
4518 name = strchr(typelist[i]->tp_name, '.');
4519 assert (name != NULL);
4520 Py_INCREF(typelist[i]);
4521 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4522 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004523
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004524 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004525}