blob: 515f3bafb32549f07853cd1c29d47cffaa9738dd [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 }
536 tdo->numread = len;
537
538 if (len == LINKCELLS) {
539 if (next != Py_None) {
540 if (Py_TYPE(next) != &teedataobject_type)
541 goto err;
542 assert(tdo->nextlink == NULL);
543 Py_INCREF(next);
544 tdo->nextlink = next;
545 }
546 } else {
547 if (next != Py_None)
548 goto err; /* shouldn't have a next if we are not full */
549 }
550 return (PyObject*)tdo;
551
552err:
553 Py_XDECREF(tdo);
554 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
555 return NULL;
556}
557
558static PyMethodDef teedataobject_methods[] = {
559 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
560 reduce_doc},
561 {NULL, NULL} /* sentinel */
562};
563
Raymond Hettingerad983e72003-11-12 14:32:26 +0000564PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000565
Raymond Hettingerad983e72003-11-12 14:32:26 +0000566static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000567 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000568 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000569 sizeof(teedataobject), /* tp_basicsize */
570 0, /* tp_itemsize */
571 /* methods */
572 (destructor)teedataobject_dealloc, /* tp_dealloc */
573 0, /* tp_print */
574 0, /* tp_getattr */
575 0, /* tp_setattr */
576 0, /* tp_reserved */
577 0, /* tp_repr */
578 0, /* tp_as_number */
579 0, /* tp_as_sequence */
580 0, /* tp_as_mapping */
581 0, /* tp_hash */
582 0, /* tp_call */
583 0, /* tp_str */
584 PyObject_GenericGetAttr, /* tp_getattro */
585 0, /* tp_setattro */
586 0, /* tp_as_buffer */
587 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
588 teedataobject_doc, /* tp_doc */
589 (traverseproc)teedataobject_traverse, /* tp_traverse */
590 (inquiry)teedataobject_clear, /* tp_clear */
591 0, /* tp_richcompare */
592 0, /* tp_weaklistoffset */
593 0, /* tp_iter */
594 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000595 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000596 0, /* tp_members */
597 0, /* tp_getset */
598 0, /* tp_base */
599 0, /* tp_dict */
600 0, /* tp_descr_get */
601 0, /* tp_descr_set */
602 0, /* tp_dictoffset */
603 0, /* tp_init */
604 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000605 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000606 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000607};
608
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000609
610static PyTypeObject tee_type;
611
612static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000613tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000614{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000615 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000616
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000617 if (to->index >= LINKCELLS) {
618 link = teedataobject_jumplink(to->dataobj);
619 Py_DECREF(to->dataobj);
620 to->dataobj = (teedataobject *)link;
621 to->index = 0;
622 }
623 value = teedataobject_getitem(to->dataobj, to->index);
624 if (value == NULL)
625 return NULL;
626 to->index++;
627 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000628}
629
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000630static int
631tee_traverse(teeobject *to, visitproc visit, void *arg)
632{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000633 Py_VISIT((PyObject *)to->dataobj);
634 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000635}
636
Raymond Hettingerad983e72003-11-12 14:32:26 +0000637static PyObject *
638tee_copy(teeobject *to)
639{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000640 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000641
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000642 newto = PyObject_GC_New(teeobject, &tee_type);
643 if (newto == NULL)
644 return NULL;
645 Py_INCREF(to->dataobj);
646 newto->dataobj = to->dataobj;
647 newto->index = to->index;
648 newto->weakreflist = NULL;
649 PyObject_GC_Track(newto);
650 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000651}
652
653PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
654
655static PyObject *
656tee_fromiterable(PyObject *iterable)
657{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000658 teeobject *to;
659 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000660
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000661 it = PyObject_GetIter(iterable);
662 if (it == NULL)
663 return NULL;
664 if (PyObject_TypeCheck(it, &tee_type)) {
665 to = (teeobject *)tee_copy((teeobject *)it);
666 goto done;
667 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000668
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000669 to = PyObject_GC_New(teeobject, &tee_type);
670 if (to == NULL)
671 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000672 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000673 if (!to->dataobj) {
674 PyObject_GC_Del(to);
675 to = NULL;
676 goto done;
677 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000679 to->index = 0;
680 to->weakreflist = NULL;
681 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000682done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000683 Py_XDECREF(it);
684 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000685}
686
687static PyObject *
688tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000690 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000691
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000692 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000693 return NULL;
694 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000695}
696
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000697static int
698tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000699{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000700 if (to->weakreflist != NULL)
701 PyObject_ClearWeakRefs((PyObject *) to);
702 Py_CLEAR(to->dataobj);
703 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000704}
705
706static void
707tee_dealloc(teeobject *to)
708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000709 PyObject_GC_UnTrack(to);
710 tee_clear(to);
711 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000712}
713
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000714static PyObject *
715tee_reduce(teeobject *to)
716{
717 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
718}
719
720static PyObject *
721tee_setstate(teeobject *to, PyObject *state)
722{
723 teedataobject *tdo;
724 int index;
725 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
726 return NULL;
727 if (index < 0 || index > LINKCELLS) {
728 PyErr_SetString(PyExc_ValueError, "Index out of range");
729 return NULL;
730 }
731 Py_CLEAR(to->dataobj);
732 to->dataobj = tdo;
733 Py_INCREF(to->dataobj);
734 to->index = index;
735 Py_RETURN_NONE;
736}
737
Raymond Hettingerad983e72003-11-12 14:32:26 +0000738PyDoc_STRVAR(teeobject_doc,
739"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000740
Raymond Hettingerad983e72003-11-12 14:32:26 +0000741static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000742 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000743 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
744 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000745 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000746};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000747
748static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000749 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000750 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000751 sizeof(teeobject), /* tp_basicsize */
752 0, /* tp_itemsize */
753 /* methods */
754 (destructor)tee_dealloc, /* tp_dealloc */
755 0, /* tp_print */
756 0, /* tp_getattr */
757 0, /* tp_setattr */
758 0, /* tp_reserved */
759 0, /* tp_repr */
760 0, /* tp_as_number */
761 0, /* tp_as_sequence */
762 0, /* tp_as_mapping */
763 0, /* tp_hash */
764 0, /* tp_call */
765 0, /* tp_str */
766 0, /* tp_getattro */
767 0, /* tp_setattro */
768 0, /* tp_as_buffer */
769 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
770 teeobject_doc, /* tp_doc */
771 (traverseproc)tee_traverse, /* tp_traverse */
772 (inquiry)tee_clear, /* tp_clear */
773 0, /* tp_richcompare */
774 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
775 PyObject_SelfIter, /* tp_iter */
776 (iternextfunc)tee_next, /* tp_iternext */
777 tee_methods, /* tp_methods */
778 0, /* tp_members */
779 0, /* tp_getset */
780 0, /* tp_base */
781 0, /* tp_dict */
782 0, /* tp_descr_get */
783 0, /* tp_descr_set */
784 0, /* tp_dictoffset */
785 0, /* tp_init */
786 0, /* tp_alloc */
787 tee_new, /* tp_new */
788 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000789};
790
Raymond Hettingerad983e72003-11-12 14:32:26 +0000791static PyObject *
792tee(PyObject *self, PyObject *args)
793{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000794 Py_ssize_t i, n=2;
795 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200796 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000797
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000798 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
799 return NULL;
800 if (n < 0) {
801 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
802 return NULL;
803 }
804 result = PyTuple_New(n);
805 if (result == NULL)
806 return NULL;
807 if (n == 0)
808 return result;
809 it = PyObject_GetIter(iterable);
810 if (it == NULL) {
811 Py_DECREF(result);
812 return NULL;
813 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200814 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000815 copyable = tee_fromiterable(it);
816 Py_DECREF(it);
817 if (copyable == NULL) {
818 Py_DECREF(result);
819 return NULL;
820 }
821 } else
822 copyable = it;
823 PyTuple_SET_ITEM(result, 0, copyable);
824 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200825
826 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000827 if (copyable == NULL) {
828 Py_DECREF(result);
829 return NULL;
830 }
831 PyTuple_SET_ITEM(result, i, copyable);
832 }
833 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000834}
835
836PyDoc_STRVAR(tee_doc,
837"tee(iterable, n=2) --> tuple of n independent iterators.");
838
839
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000840/* cycle object **********************************************************/
841
842typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000843 PyObject_HEAD
844 PyObject *it;
845 PyObject *saved;
846 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000847} cycleobject;
848
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000849static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000850
851static PyObject *
852cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
853{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000854 PyObject *it;
855 PyObject *iterable;
856 PyObject *saved;
857 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000858
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000859 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
860 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000861
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000862 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
863 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000864
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000865 /* Get iterator. */
866 it = PyObject_GetIter(iterable);
867 if (it == NULL)
868 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000869
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000870 saved = PyList_New(0);
871 if (saved == NULL) {
872 Py_DECREF(it);
873 return NULL;
874 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000875
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000876 /* create cycleobject structure */
877 lz = (cycleobject *)type->tp_alloc(type, 0);
878 if (lz == NULL) {
879 Py_DECREF(it);
880 Py_DECREF(saved);
881 return NULL;
882 }
883 lz->it = it;
884 lz->saved = saved;
885 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000886
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000887 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000888}
889
890static void
891cycle_dealloc(cycleobject *lz)
892{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000893 PyObject_GC_UnTrack(lz);
894 Py_XDECREF(lz->saved);
895 Py_XDECREF(lz->it);
896 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000897}
898
899static int
900cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
901{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000902 Py_VISIT(lz->it);
903 Py_VISIT(lz->saved);
904 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000905}
906
907static PyObject *
908cycle_next(cycleobject *lz)
909{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000910 PyObject *item;
911 PyObject *it;
912 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000913
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000914 while (1) {
915 item = PyIter_Next(lz->it);
916 if (item != NULL) {
917 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
918 Py_DECREF(item);
919 return NULL;
920 }
921 return item;
922 }
923 if (PyErr_Occurred()) {
924 if (PyErr_ExceptionMatches(PyExc_StopIteration))
925 PyErr_Clear();
926 else
927 return NULL;
928 }
929 if (PyList_Size(lz->saved) == 0)
930 return NULL;
931 it = PyObject_GetIter(lz->saved);
932 if (it == NULL)
933 return NULL;
934 tmp = lz->it;
935 lz->it = it;
936 lz->firstpass = 1;
937 Py_DECREF(tmp);
938 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000939}
940
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000941static PyObject *
942cycle_reduce(cycleobject *lz)
943{
944 /* Create a new cycle with the iterator tuple, then set
945 * the saved state on it.
946 */
947 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
948 lz->it, lz->saved, lz->firstpass);
949 }
950
951static PyObject *
952cycle_setstate(cycleobject *lz, PyObject *state)
953{
954 PyObject *saved=NULL;
955 int firstpass;
956 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
957 return NULL;
958 Py_CLEAR(lz->saved);
959 lz->saved = saved;
960 Py_XINCREF(lz->saved);
961 lz->firstpass = firstpass != 0;
962 Py_RETURN_NONE;
963}
964
965static PyMethodDef cycle_methods[] = {
966 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
967 reduce_doc},
968 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
969 setstate_doc},
970 {NULL, NULL} /* sentinel */
971};
972
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000973PyDoc_STRVAR(cycle_doc,
974"cycle(iterable) --> cycle object\n\
975\n\
976Return elements from the iterable until it is exhausted.\n\
977Then repeat the sequence indefinitely.");
978
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000979static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000980 PyVarObject_HEAD_INIT(NULL, 0)
981 "itertools.cycle", /* tp_name */
982 sizeof(cycleobject), /* tp_basicsize */
983 0, /* tp_itemsize */
984 /* methods */
985 (destructor)cycle_dealloc, /* tp_dealloc */
986 0, /* tp_print */
987 0, /* tp_getattr */
988 0, /* tp_setattr */
989 0, /* tp_reserved */
990 0, /* tp_repr */
991 0, /* tp_as_number */
992 0, /* tp_as_sequence */
993 0, /* tp_as_mapping */
994 0, /* tp_hash */
995 0, /* tp_call */
996 0, /* tp_str */
997 PyObject_GenericGetAttr, /* tp_getattro */
998 0, /* tp_setattro */
999 0, /* tp_as_buffer */
1000 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1001 Py_TPFLAGS_BASETYPE, /* tp_flags */
1002 cycle_doc, /* tp_doc */
1003 (traverseproc)cycle_traverse, /* tp_traverse */
1004 0, /* tp_clear */
1005 0, /* tp_richcompare */
1006 0, /* tp_weaklistoffset */
1007 PyObject_SelfIter, /* tp_iter */
1008 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001009 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001010 0, /* tp_members */
1011 0, /* tp_getset */
1012 0, /* tp_base */
1013 0, /* tp_dict */
1014 0, /* tp_descr_get */
1015 0, /* tp_descr_set */
1016 0, /* tp_dictoffset */
1017 0, /* tp_init */
1018 0, /* tp_alloc */
1019 cycle_new, /* tp_new */
1020 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001021};
1022
1023
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001024/* dropwhile object **********************************************************/
1025
1026typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001027 PyObject_HEAD
1028 PyObject *func;
1029 PyObject *it;
1030 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001031} dropwhileobject;
1032
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001033static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034
1035static PyObject *
1036dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1037{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001038 PyObject *func, *seq;
1039 PyObject *it;
1040 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001041
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001042 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1043 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001044
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001045 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1046 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001047
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001048 /* Get iterator. */
1049 it = PyObject_GetIter(seq);
1050 if (it == NULL)
1051 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001053 /* create dropwhileobject structure */
1054 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1055 if (lz == NULL) {
1056 Py_DECREF(it);
1057 return NULL;
1058 }
1059 Py_INCREF(func);
1060 lz->func = func;
1061 lz->it = it;
1062 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001064 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001065}
1066
1067static void
1068dropwhile_dealloc(dropwhileobject *lz)
1069{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001070 PyObject_GC_UnTrack(lz);
1071 Py_XDECREF(lz->func);
1072 Py_XDECREF(lz->it);
1073 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001074}
1075
1076static int
1077dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1078{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001079 Py_VISIT(lz->it);
1080 Py_VISIT(lz->func);
1081 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001082}
1083
1084static PyObject *
1085dropwhile_next(dropwhileobject *lz)
1086{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001087 PyObject *item, *good;
1088 PyObject *it = lz->it;
1089 long ok;
1090 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001091
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001092 iternext = *Py_TYPE(it)->tp_iternext;
1093 for (;;) {
1094 item = iternext(it);
1095 if (item == NULL)
1096 return NULL;
1097 if (lz->start == 1)
1098 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001099
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001100 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1101 if (good == NULL) {
1102 Py_DECREF(item);
1103 return NULL;
1104 }
1105 ok = PyObject_IsTrue(good);
1106 Py_DECREF(good);
1107 if (!ok) {
1108 lz->start = 1;
1109 return item;
1110 }
1111 Py_DECREF(item);
1112 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113}
1114
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001115static PyObject *
1116dropwhile_reduce(dropwhileobject *lz)
1117{
1118 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1119 lz->func, lz->it, lz->start);
1120}
1121
1122static PyObject *
1123dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1124{
1125 int start = PyObject_IsTrue(state);
1126 if (start == -1)
1127 return NULL;
1128 lz->start = start;
1129 Py_RETURN_NONE;
1130}
1131
1132static PyMethodDef dropwhile_methods[] = {
1133 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1134 reduce_doc},
1135 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1136 setstate_doc},
1137 {NULL, NULL} /* sentinel */
1138};
1139
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001140PyDoc_STRVAR(dropwhile_doc,
1141"dropwhile(predicate, iterable) --> dropwhile object\n\
1142\n\
1143Drop items from the iterable while predicate(item) is true.\n\
1144Afterwards, return every element until the iterable is exhausted.");
1145
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001146static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001147 PyVarObject_HEAD_INIT(NULL, 0)
1148 "itertools.dropwhile", /* tp_name */
1149 sizeof(dropwhileobject), /* tp_basicsize */
1150 0, /* tp_itemsize */
1151 /* methods */
1152 (destructor)dropwhile_dealloc, /* tp_dealloc */
1153 0, /* tp_print */
1154 0, /* tp_getattr */
1155 0, /* tp_setattr */
1156 0, /* tp_reserved */
1157 0, /* tp_repr */
1158 0, /* tp_as_number */
1159 0, /* tp_as_sequence */
1160 0, /* tp_as_mapping */
1161 0, /* tp_hash */
1162 0, /* tp_call */
1163 0, /* tp_str */
1164 PyObject_GenericGetAttr, /* tp_getattro */
1165 0, /* tp_setattro */
1166 0, /* tp_as_buffer */
1167 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1168 Py_TPFLAGS_BASETYPE, /* tp_flags */
1169 dropwhile_doc, /* tp_doc */
1170 (traverseproc)dropwhile_traverse, /* tp_traverse */
1171 0, /* tp_clear */
1172 0, /* tp_richcompare */
1173 0, /* tp_weaklistoffset */
1174 PyObject_SelfIter, /* tp_iter */
1175 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001176 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001177 0, /* tp_members */
1178 0, /* tp_getset */
1179 0, /* tp_base */
1180 0, /* tp_dict */
1181 0, /* tp_descr_get */
1182 0, /* tp_descr_set */
1183 0, /* tp_dictoffset */
1184 0, /* tp_init */
1185 0, /* tp_alloc */
1186 dropwhile_new, /* tp_new */
1187 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001188};
1189
1190
1191/* takewhile object **********************************************************/
1192
1193typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001194 PyObject_HEAD
1195 PyObject *func;
1196 PyObject *it;
1197 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001198} takewhileobject;
1199
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001200static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001201
1202static PyObject *
1203takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1204{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001205 PyObject *func, *seq;
1206 PyObject *it;
1207 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001208
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001209 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1210 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001211
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001212 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1213 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001214
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001215 /* Get iterator. */
1216 it = PyObject_GetIter(seq);
1217 if (it == NULL)
1218 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001219
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001220 /* create takewhileobject structure */
1221 lz = (takewhileobject *)type->tp_alloc(type, 0);
1222 if (lz == NULL) {
1223 Py_DECREF(it);
1224 return NULL;
1225 }
1226 Py_INCREF(func);
1227 lz->func = func;
1228 lz->it = it;
1229 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232}
1233
1234static void
1235takewhile_dealloc(takewhileobject *lz)
1236{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001237 PyObject_GC_UnTrack(lz);
1238 Py_XDECREF(lz->func);
1239 Py_XDECREF(lz->it);
1240 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241}
1242
1243static int
1244takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1245{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001246 Py_VISIT(lz->it);
1247 Py_VISIT(lz->func);
1248 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001249}
1250
1251static PyObject *
1252takewhile_next(takewhileobject *lz)
1253{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001254 PyObject *item, *good;
1255 PyObject *it = lz->it;
1256 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001257
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001258 if (lz->stop == 1)
1259 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001260
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001261 item = (*Py_TYPE(it)->tp_iternext)(it);
1262 if (item == NULL)
1263 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001265 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1266 if (good == NULL) {
1267 Py_DECREF(item);
1268 return NULL;
1269 }
1270 ok = PyObject_IsTrue(good);
1271 Py_DECREF(good);
1272 if (ok)
1273 return item;
1274 Py_DECREF(item);
1275 lz->stop = 1;
1276 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001277}
1278
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001279static PyObject *
1280takewhile_reduce(takewhileobject *lz)
1281{
1282 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1283 lz->func, lz->it, lz->stop);
1284}
1285
1286static PyObject *
1287takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1288{
1289 int stop = PyObject_IsTrue(state);
1290 if (stop == -1)
1291 return NULL;
1292 lz->stop = stop;
1293 Py_RETURN_NONE;
1294}
1295
1296static PyMethodDef takewhile_reduce_methods[] = {
1297 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1298 reduce_doc},
1299 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1300 setstate_doc},
1301 {NULL, NULL} /* sentinel */
1302};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001303PyDoc_STRVAR(takewhile_doc,
1304"takewhile(predicate, iterable) --> takewhile object\n\
1305\n\
1306Return successive entries from an iterable as long as the \n\
1307predicate evaluates to true for each entry.");
1308
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001309static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001310 PyVarObject_HEAD_INIT(NULL, 0)
1311 "itertools.takewhile", /* tp_name */
1312 sizeof(takewhileobject), /* tp_basicsize */
1313 0, /* tp_itemsize */
1314 /* methods */
1315 (destructor)takewhile_dealloc, /* tp_dealloc */
1316 0, /* tp_print */
1317 0, /* tp_getattr */
1318 0, /* tp_setattr */
1319 0, /* tp_reserved */
1320 0, /* tp_repr */
1321 0, /* tp_as_number */
1322 0, /* tp_as_sequence */
1323 0, /* tp_as_mapping */
1324 0, /* tp_hash */
1325 0, /* tp_call */
1326 0, /* tp_str */
1327 PyObject_GenericGetAttr, /* tp_getattro */
1328 0, /* tp_setattro */
1329 0, /* tp_as_buffer */
1330 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1331 Py_TPFLAGS_BASETYPE, /* tp_flags */
1332 takewhile_doc, /* tp_doc */
1333 (traverseproc)takewhile_traverse, /* tp_traverse */
1334 0, /* tp_clear */
1335 0, /* tp_richcompare */
1336 0, /* tp_weaklistoffset */
1337 PyObject_SelfIter, /* tp_iter */
1338 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001339 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001340 0, /* tp_members */
1341 0, /* tp_getset */
1342 0, /* tp_base */
1343 0, /* tp_dict */
1344 0, /* tp_descr_get */
1345 0, /* tp_descr_set */
1346 0, /* tp_dictoffset */
1347 0, /* tp_init */
1348 0, /* tp_alloc */
1349 takewhile_new, /* tp_new */
1350 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351};
1352
1353
1354/* islice object ************************************************************/
1355
1356typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001357 PyObject_HEAD
1358 PyObject *it;
1359 Py_ssize_t next;
1360 Py_ssize_t stop;
1361 Py_ssize_t step;
1362 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363} isliceobject;
1364
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001365static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366
1367static PyObject *
1368islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1369{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001370 PyObject *seq;
1371 Py_ssize_t start=0, stop=-1, step=1;
1372 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1373 Py_ssize_t numargs;
1374 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001376 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1377 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001378
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001379 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1380 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001382 numargs = PyTuple_Size(args);
1383 if (numargs == 2) {
1384 if (a1 != Py_None) {
1385 stop = PyLong_AsSsize_t(a1);
1386 if (stop == -1) {
1387 if (PyErr_Occurred())
1388 PyErr_Clear();
1389 PyErr_SetString(PyExc_ValueError,
1390 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1391 return NULL;
1392 }
1393 }
1394 } else {
1395 if (a1 != Py_None)
1396 start = PyLong_AsSsize_t(a1);
1397 if (start == -1 && PyErr_Occurred())
1398 PyErr_Clear();
1399 if (a2 != Py_None) {
1400 stop = PyLong_AsSsize_t(a2);
1401 if (stop == -1) {
1402 if (PyErr_Occurred())
1403 PyErr_Clear();
1404 PyErr_SetString(PyExc_ValueError,
1405 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1406 return NULL;
1407 }
1408 }
1409 }
1410 if (start<0 || stop<-1) {
1411 PyErr_SetString(PyExc_ValueError,
1412 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1413 return NULL;
1414 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001416 if (a3 != NULL) {
1417 if (a3 != Py_None)
1418 step = PyLong_AsSsize_t(a3);
1419 if (step == -1 && PyErr_Occurred())
1420 PyErr_Clear();
1421 }
1422 if (step<1) {
1423 PyErr_SetString(PyExc_ValueError,
1424 "Step for islice() must be a positive integer or None.");
1425 return NULL;
1426 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001428 /* Get iterator. */
1429 it = PyObject_GetIter(seq);
1430 if (it == NULL)
1431 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001433 /* create isliceobject structure */
1434 lz = (isliceobject *)type->tp_alloc(type, 0);
1435 if (lz == NULL) {
1436 Py_DECREF(it);
1437 return NULL;
1438 }
1439 lz->it = it;
1440 lz->next = start;
1441 lz->stop = stop;
1442 lz->step = step;
1443 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001444
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001445 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001446}
1447
1448static void
1449islice_dealloc(isliceobject *lz)
1450{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001451 PyObject_GC_UnTrack(lz);
1452 Py_XDECREF(lz->it);
1453 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454}
1455
1456static int
1457islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1458{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001459 Py_VISIT(lz->it);
1460 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461}
1462
1463static PyObject *
1464islice_next(isliceobject *lz)
1465{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001466 PyObject *item;
1467 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001468 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001469 Py_ssize_t oldnext;
1470 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001472 iternext = *Py_TYPE(it)->tp_iternext;
1473 while (lz->cnt < lz->next) {
1474 item = iternext(it);
1475 if (item == NULL)
1476 return NULL;
1477 Py_DECREF(item);
1478 lz->cnt++;
1479 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001480 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001481 return NULL;
1482 item = iternext(it);
1483 if (item == NULL)
1484 return NULL;
1485 lz->cnt++;
1486 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001487 /* The (size_t) cast below avoids the danger of undefined
1488 behaviour from signed integer overflow. */
1489 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001490 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1491 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001493}
1494
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001495static PyObject *
1496islice_reduce(isliceobject *lz)
1497{
1498 /* When unpickled, generate a new object with the same bounds,
1499 * then 'setstate' with the next and count
1500 */
1501 PyObject *stop;
1502 if (lz->stop == -1) {
1503 stop = Py_None;
1504 Py_INCREF(stop);
1505 } else {
1506 stop = PyLong_FromSsize_t(lz->stop);
1507 if (stop == NULL)
1508 return NULL;
1509 }
1510 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1511 lz->it, lz->next, stop, lz->step,
1512 lz->cnt);
1513}
1514
1515static PyObject *
1516islice_setstate(isliceobject *lz, PyObject *state)
1517{
1518 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1519 if (cnt == -1 && PyErr_Occurred())
1520 return NULL;
1521 lz->cnt = cnt;
1522 Py_RETURN_NONE;
1523}
1524
1525static PyMethodDef islice_methods[] = {
1526 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1527 reduce_doc},
1528 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1529 setstate_doc},
1530 {NULL, NULL} /* sentinel */
1531};
1532
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533PyDoc_STRVAR(islice_doc,
1534"islice(iterable, [start,] stop [, step]) --> islice object\n\
1535\n\
1536Return an iterator whose next() method returns selected values from an\n\
1537iterable. If start is specified, will skip all preceding elements;\n\
1538otherwise, start defaults to zero. Step defaults to one. If\n\
1539specified as another value, step determines how many values are \n\
1540skipped between successive calls. Works like a slice() on a list\n\
1541but returns an iterator.");
1542
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001543static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001544 PyVarObject_HEAD_INIT(NULL, 0)
1545 "itertools.islice", /* tp_name */
1546 sizeof(isliceobject), /* tp_basicsize */
1547 0, /* tp_itemsize */
1548 /* methods */
1549 (destructor)islice_dealloc, /* tp_dealloc */
1550 0, /* tp_print */
1551 0, /* tp_getattr */
1552 0, /* tp_setattr */
1553 0, /* tp_reserved */
1554 0, /* tp_repr */
1555 0, /* tp_as_number */
1556 0, /* tp_as_sequence */
1557 0, /* tp_as_mapping */
1558 0, /* tp_hash */
1559 0, /* tp_call */
1560 0, /* tp_str */
1561 PyObject_GenericGetAttr, /* tp_getattro */
1562 0, /* tp_setattro */
1563 0, /* tp_as_buffer */
1564 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1565 Py_TPFLAGS_BASETYPE, /* tp_flags */
1566 islice_doc, /* tp_doc */
1567 (traverseproc)islice_traverse, /* tp_traverse */
1568 0, /* tp_clear */
1569 0, /* tp_richcompare */
1570 0, /* tp_weaklistoffset */
1571 PyObject_SelfIter, /* tp_iter */
1572 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001573 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001574 0, /* tp_members */
1575 0, /* tp_getset */
1576 0, /* tp_base */
1577 0, /* tp_dict */
1578 0, /* tp_descr_get */
1579 0, /* tp_descr_set */
1580 0, /* tp_dictoffset */
1581 0, /* tp_init */
1582 0, /* tp_alloc */
1583 islice_new, /* tp_new */
1584 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585};
1586
1587
1588/* starmap object ************************************************************/
1589
1590typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001591 PyObject_HEAD
1592 PyObject *func;
1593 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001594} starmapobject;
1595
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001596static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001597
1598static PyObject *
1599starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1600{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001601 PyObject *func, *seq;
1602 PyObject *it;
1603 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001604
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001605 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1606 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001607
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001608 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1609 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001610
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001611 /* Get iterator. */
1612 it = PyObject_GetIter(seq);
1613 if (it == NULL)
1614 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001616 /* create starmapobject structure */
1617 lz = (starmapobject *)type->tp_alloc(type, 0);
1618 if (lz == NULL) {
1619 Py_DECREF(it);
1620 return NULL;
1621 }
1622 Py_INCREF(func);
1623 lz->func = func;
1624 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001625
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001626 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001627}
1628
1629static void
1630starmap_dealloc(starmapobject *lz)
1631{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001632 PyObject_GC_UnTrack(lz);
1633 Py_XDECREF(lz->func);
1634 Py_XDECREF(lz->it);
1635 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636}
1637
1638static int
1639starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1640{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001641 Py_VISIT(lz->it);
1642 Py_VISIT(lz->func);
1643 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644}
1645
1646static PyObject *
1647starmap_next(starmapobject *lz)
1648{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001649 PyObject *args;
1650 PyObject *result;
1651 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001652
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001653 args = (*Py_TYPE(it)->tp_iternext)(it);
1654 if (args == NULL)
1655 return NULL;
1656 if (!PyTuple_CheckExact(args)) {
1657 PyObject *newargs = PySequence_Tuple(args);
1658 Py_DECREF(args);
1659 if (newargs == NULL)
1660 return NULL;
1661 args = newargs;
1662 }
1663 result = PyObject_Call(lz->func, args, NULL);
1664 Py_DECREF(args);
1665 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666}
1667
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001668static PyObject *
1669starmap_reduce(starmapobject *lz)
1670{
1671 /* Just pickle the iterator */
1672 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1673}
1674
1675static PyMethodDef starmap_methods[] = {
1676 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1677 reduce_doc},
1678 {NULL, NULL} /* sentinel */
1679};
1680
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001681PyDoc_STRVAR(starmap_doc,
1682"starmap(function, sequence) --> starmap object\n\
1683\n\
1684Return an iterator whose values are returned from the function evaluated\n\
1685with a argument tuple taken from the given sequence.");
1686
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001687static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001688 PyVarObject_HEAD_INIT(NULL, 0)
1689 "itertools.starmap", /* tp_name */
1690 sizeof(starmapobject), /* tp_basicsize */
1691 0, /* tp_itemsize */
1692 /* methods */
1693 (destructor)starmap_dealloc, /* tp_dealloc */
1694 0, /* tp_print */
1695 0, /* tp_getattr */
1696 0, /* tp_setattr */
1697 0, /* tp_reserved */
1698 0, /* tp_repr */
1699 0, /* tp_as_number */
1700 0, /* tp_as_sequence */
1701 0, /* tp_as_mapping */
1702 0, /* tp_hash */
1703 0, /* tp_call */
1704 0, /* tp_str */
1705 PyObject_GenericGetAttr, /* tp_getattro */
1706 0, /* tp_setattro */
1707 0, /* tp_as_buffer */
1708 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1709 Py_TPFLAGS_BASETYPE, /* tp_flags */
1710 starmap_doc, /* tp_doc */
1711 (traverseproc)starmap_traverse, /* tp_traverse */
1712 0, /* tp_clear */
1713 0, /* tp_richcompare */
1714 0, /* tp_weaklistoffset */
1715 PyObject_SelfIter, /* tp_iter */
1716 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001717 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001718 0, /* tp_members */
1719 0, /* tp_getset */
1720 0, /* tp_base */
1721 0, /* tp_dict */
1722 0, /* tp_descr_get */
1723 0, /* tp_descr_set */
1724 0, /* tp_dictoffset */
1725 0, /* tp_init */
1726 0, /* tp_alloc */
1727 starmap_new, /* tp_new */
1728 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001729};
1730
1731
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001732/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001733
1734typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001735 PyObject_HEAD
1736 PyObject *source; /* Iterator over input iterables */
1737 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001738} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001740static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001741
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001742static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001743chain_new_internal(PyTypeObject *type, PyObject *source)
1744{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001745 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001746
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001747 lz = (chainobject *)type->tp_alloc(type, 0);
1748 if (lz == NULL) {
1749 Py_DECREF(source);
1750 return NULL;
1751 }
1752
1753 lz->source = source;
1754 lz->active = NULL;
1755 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001756}
1757
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001759chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001760{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001761 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001762
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001763 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1764 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001765
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001766 source = PyObject_GetIter(args);
1767 if (source == NULL)
1768 return NULL;
1769
1770 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001771}
1772
1773static PyObject *
1774chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1775{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001776 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001777
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001778 source = PyObject_GetIter(arg);
1779 if (source == NULL)
1780 return NULL;
1781
1782 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001783}
1784
1785static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001786chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001787{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001788 PyObject_GC_UnTrack(lz);
1789 Py_XDECREF(lz->active);
1790 Py_XDECREF(lz->source);
1791 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001792}
1793
Raymond Hettinger2012f172003-02-07 05:32:58 +00001794static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001795chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001796{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001797 Py_VISIT(lz->source);
1798 Py_VISIT(lz->active);
1799 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001800}
1801
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001802static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001803chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001804{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001805 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001807 if (lz->source == NULL)
1808 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001809
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001810 if (lz->active == NULL) {
1811 PyObject *iterable = PyIter_Next(lz->source);
1812 if (iterable == NULL) {
1813 Py_CLEAR(lz->source);
1814 return NULL; /* no more input sources */
1815 }
1816 lz->active = PyObject_GetIter(iterable);
1817 Py_DECREF(iterable);
1818 if (lz->active == NULL) {
1819 Py_CLEAR(lz->source);
1820 return NULL; /* input not iterable */
1821 }
1822 }
1823 item = PyIter_Next(lz->active);
1824 if (item != NULL)
1825 return item;
1826 if (PyErr_Occurred()) {
1827 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1828 PyErr_Clear();
1829 else
1830 return NULL; /* input raised an exception */
1831 }
1832 Py_CLEAR(lz->active);
1833 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001834}
1835
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001836static PyObject *
1837chain_reduce(chainobject *lz)
1838{
1839 if (lz->source) {
1840 /* we can't pickle function objects (itertools.from_iterable) so
1841 * we must use setstate to replace the iterable. One day we
1842 * will fix pickling of functions
1843 */
1844 if (lz->active) {
1845 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1846 } else {
1847 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1848 }
1849 } else {
1850 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1851 }
1852 return NULL;
1853}
1854
1855static PyObject *
1856chain_setstate(chainobject *lz, PyObject *state)
1857{
1858 PyObject *source, *active=NULL;
1859 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1860 return NULL;
1861
1862 Py_CLEAR(lz->source);
1863 lz->source = source;
1864 Py_INCREF(lz->source);
1865 Py_CLEAR(lz->active);
1866 lz->active = active;
1867 Py_XINCREF(lz->active);
1868 Py_RETURN_NONE;
1869}
1870
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001871PyDoc_STRVAR(chain_doc,
1872"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001873\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001874Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001875first iterable until it is exhausted, then elements from the next\n\
1876iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001877
Christian Heimesf16baeb2008-02-29 14:57:44 +00001878PyDoc_STRVAR(chain_from_iterable_doc,
1879"chain.from_iterable(iterable) --> chain object\n\
1880\n\
1881Alternate chain() contructor taking a single iterable argument\n\
1882that evaluates lazily.");
1883
1884static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001885 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1886 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001887 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1888 reduce_doc},
1889 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1890 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001891 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001892};
1893
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001894static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001895 PyVarObject_HEAD_INIT(NULL, 0)
1896 "itertools.chain", /* tp_name */
1897 sizeof(chainobject), /* tp_basicsize */
1898 0, /* tp_itemsize */
1899 /* methods */
1900 (destructor)chain_dealloc, /* tp_dealloc */
1901 0, /* tp_print */
1902 0, /* tp_getattr */
1903 0, /* tp_setattr */
1904 0, /* tp_reserved */
1905 0, /* tp_repr */
1906 0, /* tp_as_number */
1907 0, /* tp_as_sequence */
1908 0, /* tp_as_mapping */
1909 0, /* tp_hash */
1910 0, /* tp_call */
1911 0, /* tp_str */
1912 PyObject_GenericGetAttr, /* tp_getattro */
1913 0, /* tp_setattro */
1914 0, /* tp_as_buffer */
1915 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1916 Py_TPFLAGS_BASETYPE, /* tp_flags */
1917 chain_doc, /* tp_doc */
1918 (traverseproc)chain_traverse, /* tp_traverse */
1919 0, /* tp_clear */
1920 0, /* tp_richcompare */
1921 0, /* tp_weaklistoffset */
1922 PyObject_SelfIter, /* tp_iter */
1923 (iternextfunc)chain_next, /* tp_iternext */
1924 chain_methods, /* tp_methods */
1925 0, /* tp_members */
1926 0, /* tp_getset */
1927 0, /* tp_base */
1928 0, /* tp_dict */
1929 0, /* tp_descr_get */
1930 0, /* tp_descr_set */
1931 0, /* tp_dictoffset */
1932 0, /* tp_init */
1933 0, /* tp_alloc */
1934 chain_new, /* tp_new */
1935 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001936};
1937
1938
Christian Heimesc3f30c42008-02-22 16:37:40 +00001939/* product object ************************************************************/
1940
1941typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001942 PyObject_HEAD
1943 PyObject *pools; /* tuple of pool tuples */
1944 Py_ssize_t *indices; /* one index per pool */
1945 PyObject *result; /* most recently returned result tuple */
1946 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001947} productobject;
1948
1949static PyTypeObject product_type;
1950
1951static PyObject *
1952product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1953{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001954 productobject *lz;
1955 Py_ssize_t nargs, npools, repeat=1;
1956 PyObject *pools = NULL;
1957 Py_ssize_t *indices = NULL;
1958 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001959
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001960 if (kwds != NULL) {
1961 char *kwlist[] = {"repeat", 0};
1962 PyObject *tmpargs = PyTuple_New(0);
1963 if (tmpargs == NULL)
1964 return NULL;
1965 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1966 Py_DECREF(tmpargs);
1967 return NULL;
1968 }
1969 Py_DECREF(tmpargs);
1970 if (repeat < 0) {
1971 PyErr_SetString(PyExc_ValueError,
1972 "repeat argument cannot be negative");
1973 return NULL;
1974 }
1975 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001976
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001977 assert(PyTuple_Check(args));
1978 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1979 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001980
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001981 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1982 if (indices == NULL) {
1983 PyErr_NoMemory();
1984 goto error;
1985 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001986
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001987 pools = PyTuple_New(npools);
1988 if (pools == NULL)
1989 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001991 for (i=0; i < nargs ; ++i) {
1992 PyObject *item = PyTuple_GET_ITEM(args, i);
1993 PyObject *pool = PySequence_Tuple(item);
1994 if (pool == NULL)
1995 goto error;
1996 PyTuple_SET_ITEM(pools, i, pool);
1997 indices[i] = 0;
1998 }
1999 for ( ; i < npools; ++i) {
2000 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2001 Py_INCREF(pool);
2002 PyTuple_SET_ITEM(pools, i, pool);
2003 indices[i] = 0;
2004 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002006 /* create productobject structure */
2007 lz = (productobject *)type->tp_alloc(type, 0);
2008 if (lz == NULL)
2009 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002010
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002011 lz->pools = pools;
2012 lz->indices = indices;
2013 lz->result = NULL;
2014 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002015
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002016 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002017
2018error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002019 if (indices != NULL)
2020 PyMem_Free(indices);
2021 Py_XDECREF(pools);
2022 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002023}
2024
2025static void
2026product_dealloc(productobject *lz)
2027{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002028 PyObject_GC_UnTrack(lz);
2029 Py_XDECREF(lz->pools);
2030 Py_XDECREF(lz->result);
2031 if (lz->indices != NULL)
2032 PyMem_Free(lz->indices);
2033 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002034}
2035
2036static int
2037product_traverse(productobject *lz, visitproc visit, void *arg)
2038{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002039 Py_VISIT(lz->pools);
2040 Py_VISIT(lz->result);
2041 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002042}
2043
2044static PyObject *
2045product_next(productobject *lz)
2046{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002047 PyObject *pool;
2048 PyObject *elem;
2049 PyObject *oldelem;
2050 PyObject *pools = lz->pools;
2051 PyObject *result = lz->result;
2052 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2053 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002054
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002055 if (lz->stopped)
2056 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002057
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002058 if (result == NULL) {
2059 /* On the first pass, return an initial tuple filled with the
2060 first element from each pool. */
2061 result = PyTuple_New(npools);
2062 if (result == NULL)
2063 goto empty;
2064 lz->result = result;
2065 for (i=0; i < npools; i++) {
2066 pool = PyTuple_GET_ITEM(pools, i);
2067 if (PyTuple_GET_SIZE(pool) == 0)
2068 goto empty;
2069 elem = PyTuple_GET_ITEM(pool, 0);
2070 Py_INCREF(elem);
2071 PyTuple_SET_ITEM(result, i, elem);
2072 }
2073 } else {
2074 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002076 /* Copy the previous result tuple or re-use it if available */
2077 if (Py_REFCNT(result) > 1) {
2078 PyObject *old_result = result;
2079 result = PyTuple_New(npools);
2080 if (result == NULL)
2081 goto empty;
2082 lz->result = result;
2083 for (i=0; i < npools; i++) {
2084 elem = PyTuple_GET_ITEM(old_result, i);
2085 Py_INCREF(elem);
2086 PyTuple_SET_ITEM(result, i, elem);
2087 }
2088 Py_DECREF(old_result);
2089 }
2090 /* Now, we've got the only copy so we can update it in-place */
2091 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002093 /* Update the pool indices right-to-left. Only advance to the
2094 next pool when the previous one rolls-over */
2095 for (i=npools-1 ; i >= 0 ; i--) {
2096 pool = PyTuple_GET_ITEM(pools, i);
2097 indices[i]++;
2098 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2099 /* Roll-over and advance to next pool */
2100 indices[i] = 0;
2101 elem = PyTuple_GET_ITEM(pool, 0);
2102 Py_INCREF(elem);
2103 oldelem = PyTuple_GET_ITEM(result, i);
2104 PyTuple_SET_ITEM(result, i, elem);
2105 Py_DECREF(oldelem);
2106 } else {
2107 /* No rollover. Just increment and stop here. */
2108 elem = PyTuple_GET_ITEM(pool, indices[i]);
2109 Py_INCREF(elem);
2110 oldelem = PyTuple_GET_ITEM(result, i);
2111 PyTuple_SET_ITEM(result, i, elem);
2112 Py_DECREF(oldelem);
2113 break;
2114 }
2115 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002116
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002117 /* If i is negative, then the indices have all rolled-over
2118 and we're done. */
2119 if (i < 0)
2120 goto empty;
2121 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002123 Py_INCREF(result);
2124 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002125
2126empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002127 lz->stopped = 1;
2128 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002129}
2130
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002131static PyObject *
2132product_reduce(productobject *lz)
2133{
2134 if (lz->stopped) {
2135 return Py_BuildValue("O(())", Py_TYPE(lz));
2136 } else if (lz->result == NULL) {
2137 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2138 } else {
2139 PyObject *indices;
2140 Py_ssize_t n, i;
2141
2142 /* we must pickle the indices use them for setstate, and
2143 * additionally indicate that the iterator has started
2144 */
2145 n = PyTuple_GET_SIZE(lz->pools);
2146 indices = PyTuple_New(n);
2147 if (indices == NULL)
2148 return NULL;
2149 for (i=0; i<n; i++){
2150 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2151 if (!index) {
2152 Py_DECREF(indices);
2153 return NULL;
2154 }
2155 PyTuple_SET_ITEM(indices, i, index);
2156 }
2157 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2158 }
2159}
2160
2161static PyObject *
2162product_setstate(productobject *lz, PyObject *state)
2163{
2164 PyObject *result;
2165 Py_ssize_t n, i;
2166
2167 n = PyTuple_GET_SIZE(lz->pools);
2168 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2169 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2170 return NULL;
2171 }
2172 for (i=0; i<n; i++)
2173 {
2174 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2175 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2176 if (index < 0 && PyErr_Occurred())
2177 return NULL; /* not an integer */
2178 /* clamp the index */
2179 if (index < 0)
2180 index = 0;
2181 else if (index > n-1)
2182 index = n-1;
2183 lz->indices[i] = index;
2184 }
2185
2186 result = PyTuple_New(n);
2187 if (!result)
2188 return NULL;
2189 for (i=0; i<n; i++) {
2190 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2191 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2192 Py_INCREF(element);
2193 PyTuple_SET_ITEM(result, i, element);
2194 }
2195 Py_CLEAR(lz->result);
2196 lz->result = result;
2197 Py_RETURN_NONE;
2198}
2199
2200static PyMethodDef product_methods[] = {
2201 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2202 reduce_doc},
2203 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2204 setstate_doc},
2205 {NULL, NULL} /* sentinel */
2206};
2207
Christian Heimesc3f30c42008-02-22 16:37:40 +00002208PyDoc_STRVAR(product_doc,
2209"product(*iterables) --> product object\n\
2210\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002211Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002212For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2213The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2214cycle in a manner similar to an odometer (with the rightmost element changing\n\
2215on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002216To compute the product of an iterable with itself, specify the number\n\
2217of repetitions with the optional repeat keyword argument. For example,\n\
2218product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002219product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2220product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2221
2222static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002223 PyVarObject_HEAD_INIT(NULL, 0)
2224 "itertools.product", /* tp_name */
2225 sizeof(productobject), /* tp_basicsize */
2226 0, /* tp_itemsize */
2227 /* methods */
2228 (destructor)product_dealloc, /* tp_dealloc */
2229 0, /* tp_print */
2230 0, /* tp_getattr */
2231 0, /* tp_setattr */
2232 0, /* tp_reserved */
2233 0, /* tp_repr */
2234 0, /* tp_as_number */
2235 0, /* tp_as_sequence */
2236 0, /* tp_as_mapping */
2237 0, /* tp_hash */
2238 0, /* tp_call */
2239 0, /* tp_str */
2240 PyObject_GenericGetAttr, /* tp_getattro */
2241 0, /* tp_setattro */
2242 0, /* tp_as_buffer */
2243 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2244 Py_TPFLAGS_BASETYPE, /* tp_flags */
2245 product_doc, /* tp_doc */
2246 (traverseproc)product_traverse, /* tp_traverse */
2247 0, /* tp_clear */
2248 0, /* tp_richcompare */
2249 0, /* tp_weaklistoffset */
2250 PyObject_SelfIter, /* tp_iter */
2251 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002252 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002253 0, /* tp_members */
2254 0, /* tp_getset */
2255 0, /* tp_base */
2256 0, /* tp_dict */
2257 0, /* tp_descr_get */
2258 0, /* tp_descr_set */
2259 0, /* tp_dictoffset */
2260 0, /* tp_init */
2261 0, /* tp_alloc */
2262 product_new, /* tp_new */
2263 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002264};
2265
2266
Christian Heimes380f7f22008-02-28 11:19:05 +00002267/* combinations object ************************************************************/
2268
2269typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002270 PyObject_HEAD
2271 PyObject *pool; /* input converted to a tuple */
2272 Py_ssize_t *indices; /* one index per result element */
2273 PyObject *result; /* most recently returned result tuple */
2274 Py_ssize_t r; /* size of result tuple */
2275 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002276} combinationsobject;
2277
2278static PyTypeObject combinations_type;
2279
2280static PyObject *
2281combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002283 combinationsobject *co;
2284 Py_ssize_t n;
2285 Py_ssize_t r;
2286 PyObject *pool = NULL;
2287 PyObject *iterable = NULL;
2288 Py_ssize_t *indices = NULL;
2289 Py_ssize_t i;
2290 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002291
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002292 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2293 &iterable, &r))
2294 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002295
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002296 pool = PySequence_Tuple(iterable);
2297 if (pool == NULL)
2298 goto error;
2299 n = PyTuple_GET_SIZE(pool);
2300 if (r < 0) {
2301 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2302 goto error;
2303 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002304
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002305 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2306 if (indices == NULL) {
2307 PyErr_NoMemory();
2308 goto error;
2309 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002310
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002311 for (i=0 ; i<r ; i++)
2312 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002313
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002314 /* create combinationsobject structure */
2315 co = (combinationsobject *)type->tp_alloc(type, 0);
2316 if (co == NULL)
2317 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002318
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002319 co->pool = pool;
2320 co->indices = indices;
2321 co->result = NULL;
2322 co->r = r;
2323 co->stopped = r > n ? 1 : 0;
2324
2325 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002326
2327error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002328 if (indices != NULL)
2329 PyMem_Free(indices);
2330 Py_XDECREF(pool);
2331 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002332}
2333
2334static void
2335combinations_dealloc(combinationsobject *co)
2336{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002337 PyObject_GC_UnTrack(co);
2338 Py_XDECREF(co->pool);
2339 Py_XDECREF(co->result);
2340 if (co->indices != NULL)
2341 PyMem_Free(co->indices);
2342 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002343}
2344
2345static int
2346combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2347{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002348 Py_VISIT(co->pool);
2349 Py_VISIT(co->result);
2350 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002351}
2352
2353static PyObject *
2354combinations_next(combinationsobject *co)
2355{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002356 PyObject *elem;
2357 PyObject *oldelem;
2358 PyObject *pool = co->pool;
2359 Py_ssize_t *indices = co->indices;
2360 PyObject *result = co->result;
2361 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2362 Py_ssize_t r = co->r;
2363 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002364
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002365 if (co->stopped)
2366 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002368 if (result == NULL) {
2369 /* On the first pass, initialize result tuple using the indices */
2370 result = PyTuple_New(r);
2371 if (result == NULL)
2372 goto empty;
2373 co->result = result;
2374 for (i=0; i<r ; i++) {
2375 index = indices[i];
2376 elem = PyTuple_GET_ITEM(pool, index);
2377 Py_INCREF(elem);
2378 PyTuple_SET_ITEM(result, i, elem);
2379 }
2380 } else {
2381 /* Copy the previous result tuple or re-use it if available */
2382 if (Py_REFCNT(result) > 1) {
2383 PyObject *old_result = result;
2384 result = PyTuple_New(r);
2385 if (result == NULL)
2386 goto empty;
2387 co->result = result;
2388 for (i=0; i<r ; i++) {
2389 elem = PyTuple_GET_ITEM(old_result, i);
2390 Py_INCREF(elem);
2391 PyTuple_SET_ITEM(result, i, elem);
2392 }
2393 Py_DECREF(old_result);
2394 }
2395 /* Now, we've got the only copy so we can update it in-place
2396 * CPython's empty tuple is a singleton and cached in
2397 * PyTuple's freelist.
2398 */
2399 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002400
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002401 /* Scan indices right-to-left until finding one that is not
2402 at its maximum (i + n - r). */
2403 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2404 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002405
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002406 /* If i is negative, then the indices are all at
2407 their maximum value and we're done. */
2408 if (i < 0)
2409 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002410
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002411 /* Increment the current index which we know is not at its
2412 maximum. Then move back to the right setting each index
2413 to its lowest possible value (one higher than the index
2414 to its left -- this maintains the sort order invariant). */
2415 indices[i]++;
2416 for (j=i+1 ; j<r ; j++)
2417 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002418
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002419 /* Update the result tuple for the new indices
2420 starting with i, the leftmost index that changed */
2421 for ( ; i<r ; i++) {
2422 index = indices[i];
2423 elem = PyTuple_GET_ITEM(pool, index);
2424 Py_INCREF(elem);
2425 oldelem = PyTuple_GET_ITEM(result, i);
2426 PyTuple_SET_ITEM(result, i, elem);
2427 Py_DECREF(oldelem);
2428 }
2429 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002430
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002431 Py_INCREF(result);
2432 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002433
2434empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002435 co->stopped = 1;
2436 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002437}
2438
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002439static PyObject *
2440combinations_reduce(combinationsobject *lz)
2441{
2442 if (lz->result == NULL) {
2443 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2444 } else if (lz->stopped) {
2445 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2446 } else {
2447 PyObject *indices;
2448 Py_ssize_t i;
2449
2450 /* we must pickle the indices and use them for setstate */
2451 indices = PyTuple_New(lz->r);
2452 if (!indices)
2453 return NULL;
2454 for (i=0; i<lz->r; i++)
2455 {
2456 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2457 if (!index) {
2458 Py_DECREF(indices);
2459 return NULL;
2460 }
2461 PyTuple_SET_ITEM(indices, i, index);
2462 }
2463
2464 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2465 }
2466}
2467
2468static PyObject *
2469combinations_setstate(combinationsobject *lz, PyObject *state)
2470{
2471 PyObject *result;
2472 Py_ssize_t i;
2473 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2474
2475 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2476 {
2477 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2478 return NULL;
2479 }
2480
2481 for (i=0; i<lz->r; i++)
2482 {
2483 Py_ssize_t max;
2484 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2485 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2486 if (index == -1 && PyErr_Occurred())
2487 return NULL; /* not an integer */
2488 max = i + n - lz->r;
2489 /* clamp the index (beware of negative max) */
2490 if (index > max)
2491 index = max;
2492 if (index < 0)
2493 index = 0;
2494 lz->indices[i] = index;
2495 }
2496
2497 result = PyTuple_New(lz->r);
2498 if (result == NULL)
2499 return NULL;
2500 for (i=0; i<lz->r; i++) {
2501 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2502 Py_INCREF(element);
2503 PyTuple_SET_ITEM(result, i, element);
2504 }
2505
2506 Py_CLEAR(lz->result);
2507 lz->result = result;
2508 Py_RETURN_NONE;
2509}
2510
2511static PyMethodDef combinations_methods[] = {
2512 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2513 reduce_doc},
2514 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2515 setstate_doc},
2516 {NULL, NULL} /* sentinel */
2517};
2518
Christian Heimes380f7f22008-02-28 11:19:05 +00002519PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002520"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002521\n\
2522Return successive r-length combinations of elements in the iterable.\n\n\
2523combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2524
2525static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002526 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002527 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002528 sizeof(combinationsobject), /* tp_basicsize */
2529 0, /* tp_itemsize */
2530 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002531 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002532 0, /* tp_print */
2533 0, /* tp_getattr */
2534 0, /* tp_setattr */
2535 0, /* tp_reserved */
2536 0, /* tp_repr */
2537 0, /* tp_as_number */
2538 0, /* tp_as_sequence */
2539 0, /* tp_as_mapping */
2540 0, /* tp_hash */
2541 0, /* tp_call */
2542 0, /* tp_str */
2543 PyObject_GenericGetAttr, /* tp_getattro */
2544 0, /* tp_setattro */
2545 0, /* tp_as_buffer */
2546 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2547 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002548 combinations_doc, /* tp_doc */
2549 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002550 0, /* tp_clear */
2551 0, /* tp_richcompare */
2552 0, /* tp_weaklistoffset */
2553 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002554 (iternextfunc)combinations_next, /* tp_iternext */
2555 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002556 0, /* tp_members */
2557 0, /* tp_getset */
2558 0, /* tp_base */
2559 0, /* tp_dict */
2560 0, /* tp_descr_get */
2561 0, /* tp_descr_set */
2562 0, /* tp_dictoffset */
2563 0, /* tp_init */
2564 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002565 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002566 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002567};
2568
2569
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002570/* combinations with replacement object *******************************************/
2571
2572/* Equivalent to:
2573
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002574 def combinations_with_replacement(iterable, r):
2575 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2576 # number items returned: (n+r-1)! / r! / (n-1)!
2577 pool = tuple(iterable)
2578 n = len(pool)
2579 indices = [0] * r
2580 yield tuple(pool[i] for i in indices)
2581 while 1:
2582 for i in reversed(range(r)):
2583 if indices[i] != n - 1:
2584 break
2585 else:
2586 return
2587 indices[i:] = [indices[i] + 1] * (r - i)
2588 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002589
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002590 def combinations_with_replacement2(iterable, r):
2591 'Alternate version that filters from product()'
2592 pool = tuple(iterable)
2593 n = len(pool)
2594 for indices in product(range(n), repeat=r):
2595 if sorted(indices) == list(indices):
2596 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002597*/
2598typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002599 PyObject_HEAD
2600 PyObject *pool; /* input converted to a tuple */
2601 Py_ssize_t *indices; /* one index per result element */
2602 PyObject *result; /* most recently returned result tuple */
2603 Py_ssize_t r; /* size of result tuple */
2604 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002605} cwrobject;
2606
2607static PyTypeObject cwr_type;
2608
2609static PyObject *
2610cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2611{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002612 cwrobject *co;
2613 Py_ssize_t n;
2614 Py_ssize_t r;
2615 PyObject *pool = NULL;
2616 PyObject *iterable = NULL;
2617 Py_ssize_t *indices = NULL;
2618 Py_ssize_t i;
2619 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002620
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002621 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2622 &iterable, &r))
2623 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002624
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002625 pool = PySequence_Tuple(iterable);
2626 if (pool == NULL)
2627 goto error;
2628 n = PyTuple_GET_SIZE(pool);
2629 if (r < 0) {
2630 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2631 goto error;
2632 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002633
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002634 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2635 if (indices == NULL) {
2636 PyErr_NoMemory();
2637 goto error;
2638 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002639
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002640 for (i=0 ; i<r ; i++)
2641 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002643 /* create cwrobject structure */
2644 co = (cwrobject *)type->tp_alloc(type, 0);
2645 if (co == NULL)
2646 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002647
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002648 co->pool = pool;
2649 co->indices = indices;
2650 co->result = NULL;
2651 co->r = r;
2652 co->stopped = !n && r;
2653
2654 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002655
2656error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002657 if (indices != NULL)
2658 PyMem_Free(indices);
2659 Py_XDECREF(pool);
2660 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002661}
2662
2663static void
2664cwr_dealloc(cwrobject *co)
2665{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002666 PyObject_GC_UnTrack(co);
2667 Py_XDECREF(co->pool);
2668 Py_XDECREF(co->result);
2669 if (co->indices != NULL)
2670 PyMem_Free(co->indices);
2671 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002672}
2673
2674static int
2675cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2676{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002677 Py_VISIT(co->pool);
2678 Py_VISIT(co->result);
2679 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002680}
2681
2682static PyObject *
2683cwr_next(cwrobject *co)
2684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002685 PyObject *elem;
2686 PyObject *oldelem;
2687 PyObject *pool = co->pool;
2688 Py_ssize_t *indices = co->indices;
2689 PyObject *result = co->result;
2690 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2691 Py_ssize_t r = co->r;
2692 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002693
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002694 if (co->stopped)
2695 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002696
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002697 if (result == NULL) {
2698 /* On the first pass, initialize result tuple using the indices */
2699 result = PyTuple_New(r);
2700 if (result == NULL)
2701 goto empty;
2702 co->result = result;
2703 for (i=0; i<r ; i++) {
2704 index = indices[i];
2705 elem = PyTuple_GET_ITEM(pool, index);
2706 Py_INCREF(elem);
2707 PyTuple_SET_ITEM(result, i, elem);
2708 }
2709 } else {
2710 /* Copy the previous result tuple or re-use it if available */
2711 if (Py_REFCNT(result) > 1) {
2712 PyObject *old_result = result;
2713 result = PyTuple_New(r);
2714 if (result == NULL)
2715 goto empty;
2716 co->result = result;
2717 for (i=0; i<r ; i++) {
2718 elem = PyTuple_GET_ITEM(old_result, i);
2719 Py_INCREF(elem);
2720 PyTuple_SET_ITEM(result, i, elem);
2721 }
2722 Py_DECREF(old_result);
2723 }
2724 /* Now, we've got the only copy so we can update it in-place CPython's
2725 empty tuple is a singleton and cached in PyTuple's freelist. */
2726 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002727
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002728 /* Scan indices right-to-left until finding one that is not
2729 * at its maximum (n-1). */
2730 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2731 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002732
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002733 /* If i is negative, then the indices are all at
2734 their maximum value and we're done. */
2735 if (i < 0)
2736 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002737
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002738 /* Increment the current index which we know is not at its
2739 maximum. Then set all to the right to the same value. */
2740 indices[i]++;
2741 for (j=i+1 ; j<r ; j++)
2742 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002743
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002744 /* Update the result tuple for the new indices
2745 starting with i, the leftmost index that changed */
2746 for ( ; i<r ; i++) {
2747 index = indices[i];
2748 elem = PyTuple_GET_ITEM(pool, index);
2749 Py_INCREF(elem);
2750 oldelem = PyTuple_GET_ITEM(result, i);
2751 PyTuple_SET_ITEM(result, i, elem);
2752 Py_DECREF(oldelem);
2753 }
2754 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 Py_INCREF(result);
2757 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002758
2759empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 co->stopped = 1;
2761 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002762}
2763
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002764static PyObject *
2765cwr_reduce(cwrobject *lz)
2766{
2767 if (lz->result == NULL) {
2768 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2769 } else if (lz->stopped) {
2770 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2771 } else {
2772 PyObject *indices;
2773 Py_ssize_t i;
2774
2775 /* we must pickle the indices and use them for setstate */
2776 indices = PyTuple_New(lz->r);
2777 if (!indices)
2778 return NULL;
2779 for (i=0; i<lz->r; i++)
2780 {
2781 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2782 if (!index) {
2783 Py_DECREF(indices);
2784 return NULL;
2785 }
2786 PyTuple_SET_ITEM(indices, i, index);
2787 }
2788
2789 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2790 }
2791}
2792
2793static PyObject *
2794cwr_setstate(cwrobject *lz, PyObject *state)
2795{
2796 PyObject *result;
2797 Py_ssize_t n, i;
2798
2799 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2800 {
2801 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2802 return NULL;
2803 }
2804
2805 n = PyTuple_GET_SIZE(lz->pool);
2806 for (i=0; i<lz->r; i++)
2807 {
2808 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2809 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2810 if (index < 0 && PyErr_Occurred())
2811 return NULL; /* not an integer */
2812 /* clamp the index */
2813 if (index < 0)
2814 index = 0;
2815 else if (index > n-1)
2816 index = n-1;
2817 lz->indices[i] = index;
2818 }
2819 result = PyTuple_New(lz->r);
2820 if (result == NULL)
2821 return NULL;
2822 for (i=0; i<lz->r; i++) {
2823 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2824 Py_INCREF(element);
2825 PyTuple_SET_ITEM(result, i, element);
2826 }
2827 Py_CLEAR(lz->result);
2828 lz->result = result;
2829 Py_RETURN_NONE;
2830}
2831
2832static PyMethodDef cwr_methods[] = {
2833 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2834 reduce_doc},
2835 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2836 setstate_doc},
2837 {NULL, NULL} /* sentinel */
2838};
2839
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002840PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002841"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002842\n\
2843Return successive r-length combinations of elements in the iterable\n\
2844allowing individual elements to have successive repeats.\n\
2845combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2846
2847static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002848 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002849 "itertools.combinations_with_replacement", /* tp_name */
2850 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002851 0, /* tp_itemsize */
2852 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002853 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002854 0, /* tp_print */
2855 0, /* tp_getattr */
2856 0, /* tp_setattr */
2857 0, /* tp_reserved */
2858 0, /* tp_repr */
2859 0, /* tp_as_number */
2860 0, /* tp_as_sequence */
2861 0, /* tp_as_mapping */
2862 0, /* tp_hash */
2863 0, /* tp_call */
2864 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002865 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002866 0, /* tp_setattro */
2867 0, /* tp_as_buffer */
2868 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002869 Py_TPFLAGS_BASETYPE, /* tp_flags */
2870 cwr_doc, /* tp_doc */
2871 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002872 0, /* tp_clear */
2873 0, /* tp_richcompare */
2874 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002875 PyObject_SelfIter, /* tp_iter */
2876 (iternextfunc)cwr_next, /* tp_iternext */
2877 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002878 0, /* tp_members */
2879 0, /* tp_getset */
2880 0, /* tp_base */
2881 0, /* tp_dict */
2882 0, /* tp_descr_get */
2883 0, /* tp_descr_set */
2884 0, /* tp_dictoffset */
2885 0, /* tp_init */
2886 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002887 cwr_new, /* tp_new */
2888 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002889};
2890
2891
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002892/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002893
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002894def permutations(iterable, r=None):
2895 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2896 pool = tuple(iterable)
2897 n = len(pool)
2898 r = n if r is None else r
2899 indices = range(n)
2900 cycles = range(n-r+1, n+1)[::-1]
2901 yield tuple(pool[i] for i in indices[:r])
2902 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002903 for i in reversed(range(r)):
2904 cycles[i] -= 1
2905 if cycles[i] == 0:
2906 indices[i:] = indices[i+1:] + indices[i:i+1]
2907 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002908 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002909 j = cycles[i]
2910 indices[i], indices[-j] = indices[-j], indices[i]
2911 yield tuple(pool[i] for i in indices[:r])
2912 break
2913 else:
2914 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002915*/
2916
2917typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002918 PyObject_HEAD
2919 PyObject *pool; /* input converted to a tuple */
2920 Py_ssize_t *indices; /* one index per element in the pool */
2921 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2922 PyObject *result; /* most recently returned result tuple */
2923 Py_ssize_t r; /* size of result tuple */
2924 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002925} permutationsobject;
2926
2927static PyTypeObject permutations_type;
2928
2929static PyObject *
2930permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2931{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002932 permutationsobject *po;
2933 Py_ssize_t n;
2934 Py_ssize_t r;
2935 PyObject *robj = Py_None;
2936 PyObject *pool = NULL;
2937 PyObject *iterable = NULL;
2938 Py_ssize_t *indices = NULL;
2939 Py_ssize_t *cycles = NULL;
2940 Py_ssize_t i;
2941 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002942
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002943 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2944 &iterable, &robj))
2945 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002947 pool = PySequence_Tuple(iterable);
2948 if (pool == NULL)
2949 goto error;
2950 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002951
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002952 r = n;
2953 if (robj != Py_None) {
2954 if (!PyLong_Check(robj)) {
2955 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2956 goto error;
2957 }
2958 r = PyLong_AsSsize_t(robj);
2959 if (r == -1 && PyErr_Occurred())
2960 goto error;
2961 }
2962 if (r < 0) {
2963 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2964 goto error;
2965 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002966
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002967 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2968 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2969 if (indices == NULL || cycles == NULL) {
2970 PyErr_NoMemory();
2971 goto error;
2972 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002973
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002974 for (i=0 ; i<n ; i++)
2975 indices[i] = i;
2976 for (i=0 ; i<r ; i++)
2977 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002978
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002979 /* create permutationsobject structure */
2980 po = (permutationsobject *)type->tp_alloc(type, 0);
2981 if (po == NULL)
2982 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002983
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002984 po->pool = pool;
2985 po->indices = indices;
2986 po->cycles = cycles;
2987 po->result = NULL;
2988 po->r = r;
2989 po->stopped = r > n ? 1 : 0;
2990
2991 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002992
2993error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002994 if (indices != NULL)
2995 PyMem_Free(indices);
2996 if (cycles != NULL)
2997 PyMem_Free(cycles);
2998 Py_XDECREF(pool);
2999 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003000}
3001
3002static void
3003permutations_dealloc(permutationsobject *po)
3004{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003005 PyObject_GC_UnTrack(po);
3006 Py_XDECREF(po->pool);
3007 Py_XDECREF(po->result);
3008 PyMem_Free(po->indices);
3009 PyMem_Free(po->cycles);
3010 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003011}
3012
3013static int
3014permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3015{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003016 Py_VISIT(po->pool);
3017 Py_VISIT(po->result);
3018 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003019}
3020
3021static PyObject *
3022permutations_next(permutationsobject *po)
3023{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003024 PyObject *elem;
3025 PyObject *oldelem;
3026 PyObject *pool = po->pool;
3027 Py_ssize_t *indices = po->indices;
3028 Py_ssize_t *cycles = po->cycles;
3029 PyObject *result = po->result;
3030 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3031 Py_ssize_t r = po->r;
3032 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003033
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003034 if (po->stopped)
3035 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003036
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003037 if (result == NULL) {
3038 /* On the first pass, initialize result tuple using the indices */
3039 result = PyTuple_New(r);
3040 if (result == NULL)
3041 goto empty;
3042 po->result = result;
3043 for (i=0; i<r ; i++) {
3044 index = indices[i];
3045 elem = PyTuple_GET_ITEM(pool, index);
3046 Py_INCREF(elem);
3047 PyTuple_SET_ITEM(result, i, elem);
3048 }
3049 } else {
3050 if (n == 0)
3051 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003052
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003053 /* Copy the previous result tuple or re-use it if available */
3054 if (Py_REFCNT(result) > 1) {
3055 PyObject *old_result = result;
3056 result = PyTuple_New(r);
3057 if (result == NULL)
3058 goto empty;
3059 po->result = result;
3060 for (i=0; i<r ; i++) {
3061 elem = PyTuple_GET_ITEM(old_result, i);
3062 Py_INCREF(elem);
3063 PyTuple_SET_ITEM(result, i, elem);
3064 }
3065 Py_DECREF(old_result);
3066 }
3067 /* Now, we've got the only copy so we can update it in-place */
3068 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003069
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003070 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3071 for (i=r-1 ; i>=0 ; i--) {
3072 cycles[i] -= 1;
3073 if (cycles[i] == 0) {
3074 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3075 index = indices[i];
3076 for (j=i ; j<n-1 ; j++)
3077 indices[j] = indices[j+1];
3078 indices[n-1] = index;
3079 cycles[i] = n - i;
3080 } else {
3081 j = cycles[i];
3082 index = indices[i];
3083 indices[i] = indices[n-j];
3084 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003085
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003086 for (k=i; k<r ; k++) {
3087 /* start with i, the leftmost element that changed */
3088 /* yield tuple(pool[k] for k in indices[:r]) */
3089 index = indices[k];
3090 elem = PyTuple_GET_ITEM(pool, index);
3091 Py_INCREF(elem);
3092 oldelem = PyTuple_GET_ITEM(result, k);
3093 PyTuple_SET_ITEM(result, k, elem);
3094 Py_DECREF(oldelem);
3095 }
3096 break;
3097 }
3098 }
3099 /* If i is negative, then the cycles have all
3100 rolled-over and we're done. */
3101 if (i < 0)
3102 goto empty;
3103 }
3104 Py_INCREF(result);
3105 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003106
3107empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003108 po->stopped = 1;
3109 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003110}
3111
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003112static PyObject *
3113permutations_reduce(permutationsobject *po)
3114{
3115 if (po->result == NULL) {
3116 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3117 } else if (po->stopped) {
3118 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3119 } else {
3120 PyObject *indices=NULL, *cycles=NULL;
3121 Py_ssize_t n, i;
3122
3123 /* we must pickle the indices and cycles and use them for setstate */
3124 n = PyTuple_GET_SIZE(po->pool);
3125 indices = PyTuple_New(n);
3126 if (indices == NULL)
3127 goto err;
3128 for (i=0; i<n; i++){
3129 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3130 if (!index)
3131 goto err;
3132 PyTuple_SET_ITEM(indices, i, index);
3133 }
3134
3135 cycles = PyTuple_New(po->r);
3136 if (cycles == NULL)
3137 goto err;
3138 for (i=0; i<po->r; i++)
3139 {
3140 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3141 if (!index)
3142 goto err;
3143 PyTuple_SET_ITEM(cycles, i, index);
3144 }
3145 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3146 po->pool, po->r,
3147 indices, cycles);
3148 err:
3149 Py_XDECREF(indices);
3150 Py_XDECREF(cycles);
3151 return NULL;
3152 }
3153}
3154
3155static PyObject *
3156permutations_setstate(permutationsobject *po, PyObject *state)
3157{
3158 PyObject *indices, *cycles, *result;
3159 Py_ssize_t n, i;
3160
3161 if (!PyArg_ParseTuple(state, "O!O!",
3162 &PyTuple_Type, &indices,
3163 &PyTuple_Type, &cycles))
3164 return NULL;
3165
3166 n = PyTuple_GET_SIZE(po->pool);
3167 if (PyTuple_GET_SIZE(indices) != n ||
3168 PyTuple_GET_SIZE(cycles) != po->r)
3169 {
3170 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3171 return NULL;
3172 }
3173
3174 for (i=0; i<n; i++)
3175 {
3176 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3177 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3178 if (index < 0 && PyErr_Occurred())
3179 return NULL; /* not an integer */
3180 /* clamp the index */
3181 if (index < 0)
3182 index = 0;
3183 else if (index > n-1)
3184 index = n-1;
3185 po->indices[i] = index;
3186 }
3187
3188 for (i=0; i<po->r; i++)
3189 {
3190 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3191 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3192 if (index < 0 && PyErr_Occurred())
3193 return NULL; /* not an integer */
3194 if (index < 1)
3195 index = 1;
3196 else if (index > n-i)
3197 index = n-i;
3198 po->cycles[i] = index;
3199 }
3200 result = PyTuple_New(po->r);
3201 if (result == NULL)
3202 return NULL;
3203 for (i=0; i<po->r; i++) {
3204 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3205 Py_INCREF(element);
3206 PyTuple_SET_ITEM(result, i, element);
3207 }
3208 Py_CLEAR(po->result);
3209 po->result = result;
3210 Py_RETURN_NONE;
3211}
3212
3213static PyMethodDef permuations_methods[] = {
3214 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3215 reduce_doc},
3216 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3217 setstate_doc},
3218 {NULL, NULL} /* sentinel */
3219};
3220
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003221PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003222"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003223\n\
3224Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003225permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003226
3227static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003228 PyVarObject_HEAD_INIT(NULL, 0)
3229 "itertools.permutations", /* tp_name */
3230 sizeof(permutationsobject), /* tp_basicsize */
3231 0, /* tp_itemsize */
3232 /* methods */
3233 (destructor)permutations_dealloc, /* tp_dealloc */
3234 0, /* tp_print */
3235 0, /* tp_getattr */
3236 0, /* tp_setattr */
3237 0, /* tp_reserved */
3238 0, /* tp_repr */
3239 0, /* tp_as_number */
3240 0, /* tp_as_sequence */
3241 0, /* tp_as_mapping */
3242 0, /* tp_hash */
3243 0, /* tp_call */
3244 0, /* tp_str */
3245 PyObject_GenericGetAttr, /* tp_getattro */
3246 0, /* tp_setattro */
3247 0, /* tp_as_buffer */
3248 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3249 Py_TPFLAGS_BASETYPE, /* tp_flags */
3250 permutations_doc, /* tp_doc */
3251 (traverseproc)permutations_traverse, /* tp_traverse */
3252 0, /* tp_clear */
3253 0, /* tp_richcompare */
3254 0, /* tp_weaklistoffset */
3255 PyObject_SelfIter, /* tp_iter */
3256 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003257 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003258 0, /* tp_members */
3259 0, /* tp_getset */
3260 0, /* tp_base */
3261 0, /* tp_dict */
3262 0, /* tp_descr_get */
3263 0, /* tp_descr_set */
3264 0, /* tp_dictoffset */
3265 0, /* tp_init */
3266 0, /* tp_alloc */
3267 permutations_new, /* tp_new */
3268 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003269};
3270
Raymond Hettinger482ba772010-12-01 22:48:00 +00003271/* accumulate object ************************************************************/
3272
3273typedef struct {
3274 PyObject_HEAD
3275 PyObject *total;
3276 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003277 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003278} accumulateobject;
3279
3280static PyTypeObject accumulate_type;
3281
3282static PyObject *
3283accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3284{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003285 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003286 PyObject *iterable;
3287 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003288 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003289 accumulateobject *lz;
3290
Raymond Hettinger5d446132011-03-27 18:52:10 -07003291 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3292 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003293 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003294
3295 /* Get iterator. */
3296 it = PyObject_GetIter(iterable);
3297 if (it == NULL)
3298 return NULL;
3299
Raymond Hettinger482ba772010-12-01 22:48:00 +00003300 /* create accumulateobject structure */
3301 lz = (accumulateobject *)type->tp_alloc(type, 0);
3302 if (lz == NULL) {
3303 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003304 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003305 }
3306
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003307 if (binop != Py_None) {
3308 Py_XINCREF(binop);
3309 lz->binop = binop;
3310 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003311 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003312 lz->it = it;
3313 return (PyObject *)lz;
3314}
3315
3316static void
3317accumulate_dealloc(accumulateobject *lz)
3318{
3319 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003320 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003321 Py_XDECREF(lz->total);
3322 Py_XDECREF(lz->it);
3323 Py_TYPE(lz)->tp_free(lz);
3324}
3325
3326static int
3327accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3328{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003329 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003330 Py_VISIT(lz->it);
3331 Py_VISIT(lz->total);
3332 return 0;
3333}
3334
3335static PyObject *
3336accumulate_next(accumulateobject *lz)
3337{
3338 PyObject *val, *oldtotal, *newtotal;
3339
3340 val = PyIter_Next(lz->it);
3341 if (val == NULL)
3342 return NULL;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003343
3344 if (lz->total == NULL) {
3345 Py_INCREF(val);
3346 lz->total = val;
3347 return lz->total;
3348 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003349
3350 if (lz->binop == NULL)
3351 newtotal = PyNumber_Add(lz->total, val);
3352 else
3353 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003354 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003355 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003356 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003357
3358 oldtotal = lz->total;
3359 lz->total = newtotal;
3360 Py_DECREF(oldtotal);
3361
3362 Py_INCREF(newtotal);
3363 return newtotal;
3364}
3365
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003366static PyObject *
3367accumulate_reduce(accumulateobject *lz)
3368{
3369 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3370 lz->it, lz->binop?lz->binop:Py_None,
3371 lz->total?lz->total:Py_None);
3372 }
3373
3374static PyObject *
3375accumulate_setstate(accumulateobject *lz, PyObject *state)
3376{
3377 Py_CLEAR(lz->total);
3378 lz->total = state;
3379 Py_INCREF(lz->total);
3380 Py_RETURN_NONE;
3381}
3382
3383static PyMethodDef accumulate_methods[] = {
3384 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3385 reduce_doc},
3386 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3387 setstate_doc},
3388 {NULL, NULL} /* sentinel */
3389};
3390
Raymond Hettinger482ba772010-12-01 22:48:00 +00003391PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003392"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003393\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003394Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003395
3396static PyTypeObject accumulate_type = {
3397 PyVarObject_HEAD_INIT(NULL, 0)
3398 "itertools.accumulate", /* tp_name */
3399 sizeof(accumulateobject), /* tp_basicsize */
3400 0, /* tp_itemsize */
3401 /* methods */
3402 (destructor)accumulate_dealloc, /* tp_dealloc */
3403 0, /* tp_print */
3404 0, /* tp_getattr */
3405 0, /* tp_setattr */
3406 0, /* tp_reserved */
3407 0, /* tp_repr */
3408 0, /* tp_as_number */
3409 0, /* tp_as_sequence */
3410 0, /* tp_as_mapping */
3411 0, /* tp_hash */
3412 0, /* tp_call */
3413 0, /* tp_str */
3414 PyObject_GenericGetAttr, /* tp_getattro */
3415 0, /* tp_setattro */
3416 0, /* tp_as_buffer */
3417 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3418 Py_TPFLAGS_BASETYPE, /* tp_flags */
3419 accumulate_doc, /* tp_doc */
3420 (traverseproc)accumulate_traverse, /* tp_traverse */
3421 0, /* tp_clear */
3422 0, /* tp_richcompare */
3423 0, /* tp_weaklistoffset */
3424 PyObject_SelfIter, /* tp_iter */
3425 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003426 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003427 0, /* tp_members */
3428 0, /* tp_getset */
3429 0, /* tp_base */
3430 0, /* tp_dict */
3431 0, /* tp_descr_get */
3432 0, /* tp_descr_set */
3433 0, /* tp_dictoffset */
3434 0, /* tp_init */
3435 0, /* tp_alloc */
3436 accumulate_new, /* tp_new */
3437 PyObject_GC_Del, /* tp_free */
3438};
3439
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003440
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003441/* compress object ************************************************************/
3442
3443/* Equivalent to:
3444
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003445 def compress(data, selectors):
3446 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3447 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003448*/
3449
3450typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003451 PyObject_HEAD
3452 PyObject *data;
3453 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003454} compressobject;
3455
3456static PyTypeObject compress_type;
3457
3458static PyObject *
3459compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003461 PyObject *seq1, *seq2;
3462 PyObject *data=NULL, *selectors=NULL;
3463 compressobject *lz;
3464 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003466 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3467 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003468
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003469 data = PyObject_GetIter(seq1);
3470 if (data == NULL)
3471 goto fail;
3472 selectors = PyObject_GetIter(seq2);
3473 if (selectors == NULL)
3474 goto fail;
3475
3476 /* create compressobject structure */
3477 lz = (compressobject *)type->tp_alloc(type, 0);
3478 if (lz == NULL)
3479 goto fail;
3480 lz->data = data;
3481 lz->selectors = selectors;
3482 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003483
3484fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003485 Py_XDECREF(data);
3486 Py_XDECREF(selectors);
3487 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003488}
3489
3490static void
3491compress_dealloc(compressobject *lz)
3492{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003493 PyObject_GC_UnTrack(lz);
3494 Py_XDECREF(lz->data);
3495 Py_XDECREF(lz->selectors);
3496 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003497}
3498
3499static int
3500compress_traverse(compressobject *lz, visitproc visit, void *arg)
3501{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003502 Py_VISIT(lz->data);
3503 Py_VISIT(lz->selectors);
3504 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003505}
3506
3507static PyObject *
3508compress_next(compressobject *lz)
3509{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003510 PyObject *data = lz->data, *selectors = lz->selectors;
3511 PyObject *datum, *selector;
3512 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3513 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3514 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003515
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003516 while (1) {
3517 /* Steps: get datum, get selector, evaluate selector.
3518 Order is important (to match the pure python version
3519 in terms of which input gets a chance to raise an
3520 exception first).
3521 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003522
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003523 datum = datanext(data);
3524 if (datum == NULL)
3525 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003526
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003527 selector = selectornext(selectors);
3528 if (selector == NULL) {
3529 Py_DECREF(datum);
3530 return NULL;
3531 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003532
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003533 ok = PyObject_IsTrue(selector);
3534 Py_DECREF(selector);
3535 if (ok == 1)
3536 return datum;
3537 Py_DECREF(datum);
3538 if (ok == -1)
3539 return NULL;
3540 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003541}
3542
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003543static PyObject *
3544compress_reduce(compressobject *lz)
3545{
3546 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3547 lz->data, lz->selectors);
3548 }
3549
3550static PyMethodDef compress_methods[] = {
3551 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3552 reduce_doc},
3553 {NULL, NULL} /* sentinel */
3554};
3555
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003556PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003557"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003558\n\
3559Return data elements corresponding to true selector elements.\n\
3560Forms a shorter iterator from selected data elements using the\n\
3561selectors to choose the data elements.");
3562
3563static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003564 PyVarObject_HEAD_INIT(NULL, 0)
3565 "itertools.compress", /* tp_name */
3566 sizeof(compressobject), /* tp_basicsize */
3567 0, /* tp_itemsize */
3568 /* methods */
3569 (destructor)compress_dealloc, /* tp_dealloc */
3570 0, /* tp_print */
3571 0, /* tp_getattr */
3572 0, /* tp_setattr */
3573 0, /* tp_reserved */
3574 0, /* tp_repr */
3575 0, /* tp_as_number */
3576 0, /* tp_as_sequence */
3577 0, /* tp_as_mapping */
3578 0, /* tp_hash */
3579 0, /* tp_call */
3580 0, /* tp_str */
3581 PyObject_GenericGetAttr, /* tp_getattro */
3582 0, /* tp_setattro */
3583 0, /* tp_as_buffer */
3584 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3585 Py_TPFLAGS_BASETYPE, /* tp_flags */
3586 compress_doc, /* tp_doc */
3587 (traverseproc)compress_traverse, /* tp_traverse */
3588 0, /* tp_clear */
3589 0, /* tp_richcompare */
3590 0, /* tp_weaklistoffset */
3591 PyObject_SelfIter, /* tp_iter */
3592 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003593 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003594 0, /* tp_members */
3595 0, /* tp_getset */
3596 0, /* tp_base */
3597 0, /* tp_dict */
3598 0, /* tp_descr_get */
3599 0, /* tp_descr_set */
3600 0, /* tp_dictoffset */
3601 0, /* tp_init */
3602 0, /* tp_alloc */
3603 compress_new, /* tp_new */
3604 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003605};
3606
3607
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003608/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003609
3610typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003611 PyObject_HEAD
3612 PyObject *func;
3613 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003614} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003615
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003616static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003617
3618static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003619filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003620{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003621 PyObject *func, *seq;
3622 PyObject *it;
3623 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003624
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003625 if (type == &filterfalse_type &&
3626 !_PyArg_NoKeywords("filterfalse()", kwds))
3627 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003628
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003629 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3630 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003631
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003632 /* Get iterator. */
3633 it = PyObject_GetIter(seq);
3634 if (it == NULL)
3635 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003636
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003637 /* create filterfalseobject structure */
3638 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3639 if (lz == NULL) {
3640 Py_DECREF(it);
3641 return NULL;
3642 }
3643 Py_INCREF(func);
3644 lz->func = func;
3645 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003646
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003647 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003648}
3649
3650static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003651filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003652{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003653 PyObject_GC_UnTrack(lz);
3654 Py_XDECREF(lz->func);
3655 Py_XDECREF(lz->it);
3656 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003657}
3658
3659static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003660filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003661{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003662 Py_VISIT(lz->it);
3663 Py_VISIT(lz->func);
3664 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003665}
3666
3667static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003668filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003669{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003670 PyObject *item;
3671 PyObject *it = lz->it;
3672 long ok;
3673 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003675 iternext = *Py_TYPE(it)->tp_iternext;
3676 for (;;) {
3677 item = iternext(it);
3678 if (item == NULL)
3679 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003680
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003681 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3682 ok = PyObject_IsTrue(item);
3683 } else {
3684 PyObject *good;
3685 good = PyObject_CallFunctionObjArgs(lz->func,
3686 item, NULL);
3687 if (good == NULL) {
3688 Py_DECREF(item);
3689 return NULL;
3690 }
3691 ok = PyObject_IsTrue(good);
3692 Py_DECREF(good);
3693 }
3694 if (!ok)
3695 return item;
3696 Py_DECREF(item);
3697 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003698}
3699
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003700static PyObject *
3701filterfalse_reduce(filterfalseobject *lz)
3702{
3703 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3704 lz->func, lz->it);
3705 }
3706
3707static PyMethodDef filterfalse_methods[] = {
3708 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3709 reduce_doc},
3710 {NULL, NULL} /* sentinel */
3711};
3712
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003713PyDoc_STRVAR(filterfalse_doc,
3714"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003715\n\
3716Return those items of sequence for which function(item) is false.\n\
3717If function is None, return the items that are false.");
3718
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003719static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003720 PyVarObject_HEAD_INIT(NULL, 0)
3721 "itertools.filterfalse", /* tp_name */
3722 sizeof(filterfalseobject), /* tp_basicsize */
3723 0, /* tp_itemsize */
3724 /* methods */
3725 (destructor)filterfalse_dealloc, /* tp_dealloc */
3726 0, /* tp_print */
3727 0, /* tp_getattr */
3728 0, /* tp_setattr */
3729 0, /* tp_reserved */
3730 0, /* tp_repr */
3731 0, /* tp_as_number */
3732 0, /* tp_as_sequence */
3733 0, /* tp_as_mapping */
3734 0, /* tp_hash */
3735 0, /* tp_call */
3736 0, /* tp_str */
3737 PyObject_GenericGetAttr, /* tp_getattro */
3738 0, /* tp_setattro */
3739 0, /* tp_as_buffer */
3740 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3741 Py_TPFLAGS_BASETYPE, /* tp_flags */
3742 filterfalse_doc, /* tp_doc */
3743 (traverseproc)filterfalse_traverse, /* tp_traverse */
3744 0, /* tp_clear */
3745 0, /* tp_richcompare */
3746 0, /* tp_weaklistoffset */
3747 PyObject_SelfIter, /* tp_iter */
3748 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003749 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003750 0, /* tp_members */
3751 0, /* tp_getset */
3752 0, /* tp_base */
3753 0, /* tp_dict */
3754 0, /* tp_descr_get */
3755 0, /* tp_descr_set */
3756 0, /* tp_dictoffset */
3757 0, /* tp_init */
3758 0, /* tp_alloc */
3759 filterfalse_new, /* tp_new */
3760 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003761};
3762
3763
3764/* count object ************************************************************/
3765
3766typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003767 PyObject_HEAD
3768 Py_ssize_t cnt;
3769 PyObject *long_cnt;
3770 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003771} countobject;
3772
Raymond Hettinger30729212009-02-12 06:28:27 +00003773/* Counting logic and invariants:
3774
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003775fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003776
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003777 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3778 Advances with: cnt += 1
3779 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003780
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003781slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
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 != NULL);
3784 All counting is done with python objects (no overflows or underflows).
3785 Advances with: long_cnt += long_step
3786 Step may be zero -- effectively a slow version of repeat(cnt).
3787 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003788*/
3789
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003790static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003791
3792static PyObject *
3793count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3794{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003795 countobject *lz;
3796 int slow_mode = 0;
3797 Py_ssize_t cnt = 0;
3798 PyObject *long_cnt = NULL;
3799 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003800 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003801 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003802
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003803 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3804 kwlist, &long_cnt, &long_step))
3805 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003806
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003807 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3808 (long_step != NULL && !PyNumber_Check(long_step))) {
3809 PyErr_SetString(PyExc_TypeError, "a number is required");
3810 return NULL;
3811 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003813 if (long_cnt != NULL) {
3814 cnt = PyLong_AsSsize_t(long_cnt);
3815 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3816 PyErr_Clear();
3817 slow_mode = 1;
3818 }
3819 Py_INCREF(long_cnt);
3820 } else {
3821 cnt = 0;
3822 long_cnt = PyLong_FromLong(0);
3823 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003824
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003825 /* If not specified, step defaults to 1 */
3826 if (long_step == NULL) {
3827 long_step = PyLong_FromLong(1);
3828 if (long_step == NULL) {
3829 Py_DECREF(long_cnt);
3830 return NULL;
3831 }
3832 } else
3833 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003834
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003835 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003836
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003837 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003838 step = PyLong_AsLong(long_step);
3839 if (step != 1) {
3840 slow_mode = 1;
3841 if (step == -1 && PyErr_Occurred())
3842 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003843 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003844
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003845 if (slow_mode)
3846 cnt = PY_SSIZE_T_MAX;
3847 else
3848 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003850 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3851 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3852 assert(slow_mode ||
3853 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003854
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003855 /* create countobject structure */
3856 lz = (countobject *)type->tp_alloc(type, 0);
3857 if (lz == NULL) {
3858 Py_XDECREF(long_cnt);
3859 return NULL;
3860 }
3861 lz->cnt = cnt;
3862 lz->long_cnt = long_cnt;
3863 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003864
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003865 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003866}
3867
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003868static void
3869count_dealloc(countobject *lz)
3870{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003871 PyObject_GC_UnTrack(lz);
3872 Py_XDECREF(lz->long_cnt);
3873 Py_XDECREF(lz->long_step);
3874 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003875}
3876
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003877static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003878count_traverse(countobject *lz, visitproc visit, void *arg)
3879{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003880 Py_VISIT(lz->long_cnt);
3881 Py_VISIT(lz->long_step);
3882 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003883}
3884
3885static PyObject *
3886count_nextlong(countobject *lz)
3887{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003888 PyObject *long_cnt;
3889 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003891 long_cnt = lz->long_cnt;
3892 if (long_cnt == NULL) {
3893 /* Switch to slow_mode */
3894 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3895 if (long_cnt == NULL)
3896 return NULL;
3897 }
3898 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003899
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003900 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3901 if (stepped_up == NULL)
3902 return NULL;
3903 lz->long_cnt = stepped_up;
3904 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003905}
3906
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003907static PyObject *
3908count_next(countobject *lz)
3909{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003910 if (lz->cnt == PY_SSIZE_T_MAX)
3911 return count_nextlong(lz);
3912 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003913}
3914
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003915static PyObject *
3916count_repr(countobject *lz)
3917{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003918 if (lz->cnt != PY_SSIZE_T_MAX)
3919 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003920
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003921 if (PyLong_Check(lz->long_step)) {
3922 long step = PyLong_AsLong(lz->long_step);
3923 if (step == -1 && PyErr_Occurred()) {
3924 PyErr_Clear();
3925 }
3926 if (step == 1) {
3927 /* Don't display step when it is an integer equal to 1 */
3928 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3929 }
3930 }
3931 return PyUnicode_FromFormat("count(%R, %R)",
3932 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003933}
3934
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003935static PyObject *
3936count_reduce(countobject *lz)
3937{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003938 if (lz->cnt == PY_SSIZE_T_MAX)
3939 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3940 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003941}
3942
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003943static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003944 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003945 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003946 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003947};
3948
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003949PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003950 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003951\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003952Return a count object whose .__next__() method returns consecutive values.\n\
3953Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003954 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003955 x = firstval\n\
3956 while 1:\n\
3957 yield x\n\
3958 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003959
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003960static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003961 PyVarObject_HEAD_INIT(NULL, 0)
3962 "itertools.count", /* tp_name */
3963 sizeof(countobject), /* tp_basicsize */
3964 0, /* tp_itemsize */
3965 /* methods */
3966 (destructor)count_dealloc, /* tp_dealloc */
3967 0, /* tp_print */
3968 0, /* tp_getattr */
3969 0, /* tp_setattr */
3970 0, /* tp_reserved */
3971 (reprfunc)count_repr, /* tp_repr */
3972 0, /* tp_as_number */
3973 0, /* tp_as_sequence */
3974 0, /* tp_as_mapping */
3975 0, /* tp_hash */
3976 0, /* tp_call */
3977 0, /* tp_str */
3978 PyObject_GenericGetAttr, /* tp_getattro */
3979 0, /* tp_setattro */
3980 0, /* tp_as_buffer */
3981 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3982 Py_TPFLAGS_BASETYPE, /* tp_flags */
3983 count_doc, /* tp_doc */
3984 (traverseproc)count_traverse, /* tp_traverse */
3985 0, /* tp_clear */
3986 0, /* tp_richcompare */
3987 0, /* tp_weaklistoffset */
3988 PyObject_SelfIter, /* tp_iter */
3989 (iternextfunc)count_next, /* tp_iternext */
3990 count_methods, /* tp_methods */
3991 0, /* tp_members */
3992 0, /* tp_getset */
3993 0, /* tp_base */
3994 0, /* tp_dict */
3995 0, /* tp_descr_get */
3996 0, /* tp_descr_set */
3997 0, /* tp_dictoffset */
3998 0, /* tp_init */
3999 0, /* tp_alloc */
4000 count_new, /* tp_new */
4001 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004002};
4003
4004
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004005/* repeat object ************************************************************/
4006
4007typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004008 PyObject_HEAD
4009 PyObject *element;
4010 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004011} repeatobject;
4012
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004013static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004014
4015static PyObject *
4016repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4017{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004018 repeatobject *ro;
4019 PyObject *element;
4020 Py_ssize_t cnt = -1;
4021 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004022
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004023 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4024 &element, &cnt))
4025 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004026
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004027 if (PyTuple_Size(args) == 2 && cnt < 0)
4028 cnt = 0;
4029
4030 ro = (repeatobject *)type->tp_alloc(type, 0);
4031 if (ro == NULL)
4032 return NULL;
4033 Py_INCREF(element);
4034 ro->element = element;
4035 ro->cnt = cnt;
4036 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004037}
4038
4039static void
4040repeat_dealloc(repeatobject *ro)
4041{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004042 PyObject_GC_UnTrack(ro);
4043 Py_XDECREF(ro->element);
4044 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004045}
4046
4047static int
4048repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4049{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004050 Py_VISIT(ro->element);
4051 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004052}
4053
4054static PyObject *
4055repeat_next(repeatobject *ro)
4056{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004057 if (ro->cnt == 0)
4058 return NULL;
4059 if (ro->cnt > 0)
4060 ro->cnt--;
4061 Py_INCREF(ro->element);
4062 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004063}
4064
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004065static PyObject *
4066repeat_repr(repeatobject *ro)
4067{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004068 if (ro->cnt == -1)
4069 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4070 else
4071 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4072}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004073
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004074static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004075repeat_len(repeatobject *ro)
4076{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004077 if (ro->cnt == -1) {
4078 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4079 return NULL;
4080 }
4081 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004082}
4083
Armin Rigof5b3e362006-02-11 21:32:43 +00004084PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004085
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004086static PyObject *
4087repeat_reduce(repeatobject *ro)
4088{
4089 /* unpickle this so that a new repeat iterator is constructed with an
4090 * object, then call __setstate__ on it to set cnt
4091 */
4092 if (ro->cnt >= 0)
4093 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4094 else
4095 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4096}
4097
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004098static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004099 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004100 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004101 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004102};
4103
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004104PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004105"repeat(object [,times]) -> create an iterator which returns the object\n\
4106for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004107endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004108
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004109static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004110 PyVarObject_HEAD_INIT(NULL, 0)
4111 "itertools.repeat", /* tp_name */
4112 sizeof(repeatobject), /* tp_basicsize */
4113 0, /* tp_itemsize */
4114 /* methods */
4115 (destructor)repeat_dealloc, /* tp_dealloc */
4116 0, /* tp_print */
4117 0, /* tp_getattr */
4118 0, /* tp_setattr */
4119 0, /* tp_reserved */
4120 (reprfunc)repeat_repr, /* tp_repr */
4121 0, /* tp_as_number */
4122 0, /* tp_as_sequence */
4123 0, /* tp_as_mapping */
4124 0, /* tp_hash */
4125 0, /* tp_call */
4126 0, /* tp_str */
4127 PyObject_GenericGetAttr, /* tp_getattro */
4128 0, /* tp_setattro */
4129 0, /* tp_as_buffer */
4130 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4131 Py_TPFLAGS_BASETYPE, /* tp_flags */
4132 repeat_doc, /* tp_doc */
4133 (traverseproc)repeat_traverse, /* tp_traverse */
4134 0, /* tp_clear */
4135 0, /* tp_richcompare */
4136 0, /* tp_weaklistoffset */
4137 PyObject_SelfIter, /* tp_iter */
4138 (iternextfunc)repeat_next, /* tp_iternext */
4139 repeat_methods, /* tp_methods */
4140 0, /* tp_members */
4141 0, /* tp_getset */
4142 0, /* tp_base */
4143 0, /* tp_dict */
4144 0, /* tp_descr_get */
4145 0, /* tp_descr_set */
4146 0, /* tp_dictoffset */
4147 0, /* tp_init */
4148 0, /* tp_alloc */
4149 repeat_new, /* tp_new */
4150 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004151};
4152
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004153/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004154
4155#include "Python.h"
4156
4157typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004158 PyObject_HEAD
4159 Py_ssize_t tuplesize;
4160 Py_ssize_t numactive;
4161 PyObject *ittuple; /* tuple of iterators */
4162 PyObject *result;
4163 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004164} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004165
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004166static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004167
4168static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004169zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004170{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004171 ziplongestobject *lz;
4172 Py_ssize_t i;
4173 PyObject *ittuple; /* tuple of iterators */
4174 PyObject *result;
4175 PyObject *fillvalue = Py_None;
4176 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004178 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4179 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4180 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4181 PyErr_SetString(PyExc_TypeError,
4182 "zip_longest() got an unexpected keyword argument");
4183 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004184 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004185 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004186
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004187 /* args must be a tuple */
4188 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004189
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004190 /* obtain iterators */
4191 ittuple = PyTuple_New(tuplesize);
4192 if (ittuple == NULL)
4193 return NULL;
4194 for (i=0; i < tuplesize; ++i) {
4195 PyObject *item = PyTuple_GET_ITEM(args, i);
4196 PyObject *it = PyObject_GetIter(item);
4197 if (it == NULL) {
4198 if (PyErr_ExceptionMatches(PyExc_TypeError))
4199 PyErr_Format(PyExc_TypeError,
4200 "zip_longest argument #%zd must support iteration",
4201 i+1);
4202 Py_DECREF(ittuple);
4203 return NULL;
4204 }
4205 PyTuple_SET_ITEM(ittuple, i, it);
4206 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004207
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004208 /* create a result holder */
4209 result = PyTuple_New(tuplesize);
4210 if (result == NULL) {
4211 Py_DECREF(ittuple);
4212 return NULL;
4213 }
4214 for (i=0 ; i < tuplesize ; i++) {
4215 Py_INCREF(Py_None);
4216 PyTuple_SET_ITEM(result, i, Py_None);
4217 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004218
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004219 /* create ziplongestobject structure */
4220 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4221 if (lz == NULL) {
4222 Py_DECREF(ittuple);
4223 Py_DECREF(result);
4224 return NULL;
4225 }
4226 lz->ittuple = ittuple;
4227 lz->tuplesize = tuplesize;
4228 lz->numactive = tuplesize;
4229 lz->result = result;
4230 Py_INCREF(fillvalue);
4231 lz->fillvalue = fillvalue;
4232 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004233}
4234
4235static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004236zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004237{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004238 PyObject_GC_UnTrack(lz);
4239 Py_XDECREF(lz->ittuple);
4240 Py_XDECREF(lz->result);
4241 Py_XDECREF(lz->fillvalue);
4242 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004243}
4244
4245static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004246zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004247{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004248 Py_VISIT(lz->ittuple);
4249 Py_VISIT(lz->result);
4250 Py_VISIT(lz->fillvalue);
4251 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004252}
4253
4254static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004255zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004257 Py_ssize_t i;
4258 Py_ssize_t tuplesize = lz->tuplesize;
4259 PyObject *result = lz->result;
4260 PyObject *it;
4261 PyObject *item;
4262 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004263
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004264 if (tuplesize == 0)
4265 return NULL;
4266 if (lz->numactive == 0)
4267 return NULL;
4268 if (Py_REFCNT(result) == 1) {
4269 Py_INCREF(result);
4270 for (i=0 ; i < tuplesize ; i++) {
4271 it = PyTuple_GET_ITEM(lz->ittuple, i);
4272 if (it == NULL) {
4273 Py_INCREF(lz->fillvalue);
4274 item = lz->fillvalue;
4275 } else {
4276 item = PyIter_Next(it);
4277 if (item == NULL) {
4278 lz->numactive -= 1;
4279 if (lz->numactive == 0 || PyErr_Occurred()) {
4280 lz->numactive = 0;
4281 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004282 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004283 } else {
4284 Py_INCREF(lz->fillvalue);
4285 item = lz->fillvalue;
4286 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4287 Py_DECREF(it);
4288 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004289 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004290 }
4291 olditem = PyTuple_GET_ITEM(result, i);
4292 PyTuple_SET_ITEM(result, i, item);
4293 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004294 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004295 } else {
4296 result = PyTuple_New(tuplesize);
4297 if (result == NULL)
4298 return NULL;
4299 for (i=0 ; i < tuplesize ; i++) {
4300 it = PyTuple_GET_ITEM(lz->ittuple, i);
4301 if (it == NULL) {
4302 Py_INCREF(lz->fillvalue);
4303 item = lz->fillvalue;
4304 } else {
4305 item = PyIter_Next(it);
4306 if (item == NULL) {
4307 lz->numactive -= 1;
4308 if (lz->numactive == 0 || PyErr_Occurred()) {
4309 lz->numactive = 0;
4310 Py_DECREF(result);
4311 return NULL;
4312 } else {
4313 Py_INCREF(lz->fillvalue);
4314 item = lz->fillvalue;
4315 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4316 Py_DECREF(it);
4317 }
4318 }
4319 }
4320 PyTuple_SET_ITEM(result, i, item);
4321 }
4322 }
4323 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004324}
4325
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004326static PyObject *
4327zip_longest_reduce(ziplongestobject *lz)
4328{
4329
4330 /* Create a new tuple with empty sequences where appropriate to pickle.
4331 * Then use setstate to set the fillvalue
4332 */
4333 int i;
4334 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4335 if (args == NULL)
4336 return NULL;
4337 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4338 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4339 if (elem == NULL) {
4340 elem = PyTuple_New(0);
4341 if (elem == NULL) {
4342 Py_DECREF(args);
4343 return NULL;
4344 }
4345 } else
4346 Py_INCREF(elem);
4347 PyTuple_SET_ITEM(args, i, elem);
4348 }
4349 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4350}
4351
4352static PyObject *
4353zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4354{
4355 Py_CLEAR(lz->fillvalue);
4356 lz->fillvalue = state;
4357 Py_INCREF(lz->fillvalue);
4358 Py_RETURN_NONE;
4359}
4360
4361static PyMethodDef zip_longest_methods[] = {
4362 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4363 reduce_doc},
4364 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4365 setstate_doc},
4366 {NULL, NULL} /* sentinel */
4367};
4368
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004369PyDoc_STRVAR(zip_longest_doc,
4370"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004371\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004372Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004373the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004374method continues until the longest iterable in the argument sequence\n\
4375is exhausted and then it raises StopIteration. When the shorter iterables\n\
4376are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4377defaults to None or can be specified by a keyword argument.\n\
4378");
4379
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004380static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004381 PyVarObject_HEAD_INIT(NULL, 0)
4382 "itertools.zip_longest", /* tp_name */
4383 sizeof(ziplongestobject), /* tp_basicsize */
4384 0, /* tp_itemsize */
4385 /* methods */
4386 (destructor)zip_longest_dealloc, /* tp_dealloc */
4387 0, /* tp_print */
4388 0, /* tp_getattr */
4389 0, /* tp_setattr */
4390 0, /* tp_reserved */
4391 0, /* tp_repr */
4392 0, /* tp_as_number */
4393 0, /* tp_as_sequence */
4394 0, /* tp_as_mapping */
4395 0, /* tp_hash */
4396 0, /* tp_call */
4397 0, /* tp_str */
4398 PyObject_GenericGetAttr, /* tp_getattro */
4399 0, /* tp_setattro */
4400 0, /* tp_as_buffer */
4401 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4402 Py_TPFLAGS_BASETYPE, /* tp_flags */
4403 zip_longest_doc, /* tp_doc */
4404 (traverseproc)zip_longest_traverse, /* tp_traverse */
4405 0, /* tp_clear */
4406 0, /* tp_richcompare */
4407 0, /* tp_weaklistoffset */
4408 PyObject_SelfIter, /* tp_iter */
4409 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004410 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004411 0, /* tp_members */
4412 0, /* tp_getset */
4413 0, /* tp_base */
4414 0, /* tp_dict */
4415 0, /* tp_descr_get */
4416 0, /* tp_descr_set */
4417 0, /* tp_dictoffset */
4418 0, /* tp_init */
4419 0, /* tp_alloc */
4420 zip_longest_new, /* tp_new */
4421 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004422};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004423
4424/* module level code ********************************************************/
4425
4426PyDoc_STRVAR(module_doc,
4427"Functional tools for creating and using iterators.\n\
4428\n\
4429Infinite iterators:\n\
4430count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004431cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004432repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004433\n\
4434Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004435accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004436chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
4437compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4438dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4439groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004440filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004441islice(seq, [start,] stop [, step]) --> elements from\n\
4442 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004443starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004444tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004445takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004446zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004447\n\
4448Combinatoric generators:\n\
4449product(p, q, ... [repeat=1]) --> cartesian product\n\
4450permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004451combinations(p, r)\n\
4452combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004453");
4454
4455
Raymond Hettingerad983e72003-11-12 14:32:26 +00004456static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004457 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4458 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004459};
4460
Martin v. Löwis1a214512008-06-11 05:26:20 +00004461
4462static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004463 PyModuleDef_HEAD_INIT,
4464 "itertools",
4465 module_doc,
4466 -1,
4467 module_methods,
4468 NULL,
4469 NULL,
4470 NULL,
4471 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004472};
4473
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004474PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004475PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004476{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004477 int i;
4478 PyObject *m;
4479 char *name;
4480 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004481 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004482 &combinations_type,
4483 &cwr_type,
4484 &cycle_type,
4485 &dropwhile_type,
4486 &takewhile_type,
4487 &islice_type,
4488 &starmap_type,
4489 &chain_type,
4490 &compress_type,
4491 &filterfalse_type,
4492 &count_type,
4493 &ziplongest_type,
4494 &permutations_type,
4495 &product_type,
4496 &repeat_type,
4497 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004498 &_grouper_type,
4499 &tee_type,
4500 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004501 NULL
4502 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004503
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004504 Py_TYPE(&teedataobject_type) = &PyType_Type;
4505 m = PyModule_Create(&itertoolsmodule);
4506 if (m == NULL)
4507 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004508
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004509 for (i=0 ; typelist[i] != NULL ; i++) {
4510 if (PyType_Ready(typelist[i]) < 0)
4511 return NULL;
4512 name = strchr(typelist[i]->tp_name, '.');
4513 assert (name != NULL);
4514 Py_INCREF(typelist[i]);
4515 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4516 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004517
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004518 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004519}