blob: 194f7fb5a39f116432ab1288975cb9c857584188 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00005/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000015 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000021} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000029 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000032
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000056 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000074}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000081 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000089
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000090 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000137static PyObject *
138groupby_reduce(groupbyobject *lz)
139{
140 /* reduce as a 'new' call with an optional 'setstate' if groupby
141 * has started
142 */
143 PyObject *value;
144 if (lz->tgtkey && lz->currkey && lz->currvalue)
145 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
146 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
147 else
148 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
149 lz->it, lz->keyfunc);
150
151 return value;
152}
153
154PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
155
156static PyObject *
157groupby_setstate(groupbyobject *lz, PyObject *state)
158{
159 PyObject *currkey, *currvalue, *tgtkey;
160 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
161 return NULL;
162 Py_CLEAR(lz->currkey);
163 lz->currkey = currkey;
164 Py_INCREF(lz->currkey);
165 Py_CLEAR(lz->currvalue);
166 lz->currvalue = currvalue;
167 Py_INCREF(lz->currvalue);
168 Py_CLEAR(lz->tgtkey);
169 lz->tgtkey = tgtkey;
170 Py_INCREF(lz->tgtkey);
171 Py_RETURN_NONE;
172}
173
174PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
175
176static PyMethodDef groupby_methods[] = {
177 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
178 reduce_doc},
179 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
180 setstate_doc},
181 {NULL, NULL} /* sentinel */
182};
183
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000184PyDoc_STRVAR(groupby_doc,
185"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
186(key, sub-iterator) grouped by each value of key(value).\n");
187
188static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000189 PyVarObject_HEAD_INIT(NULL, 0)
190 "itertools.groupby", /* tp_name */
191 sizeof(groupbyobject), /* tp_basicsize */
192 0, /* tp_itemsize */
193 /* methods */
194 (destructor)groupby_dealloc, /* tp_dealloc */
195 0, /* tp_print */
196 0, /* tp_getattr */
197 0, /* tp_setattr */
198 0, /* tp_reserved */
199 0, /* tp_repr */
200 0, /* tp_as_number */
201 0, /* tp_as_sequence */
202 0, /* tp_as_mapping */
203 0, /* tp_hash */
204 0, /* tp_call */
205 0, /* tp_str */
206 PyObject_GenericGetAttr, /* tp_getattro */
207 0, /* tp_setattro */
208 0, /* tp_as_buffer */
209 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
210 Py_TPFLAGS_BASETYPE, /* tp_flags */
211 groupby_doc, /* tp_doc */
212 (traverseproc)groupby_traverse, /* tp_traverse */
213 0, /* tp_clear */
214 0, /* tp_richcompare */
215 0, /* tp_weaklistoffset */
216 PyObject_SelfIter, /* tp_iter */
217 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000218 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000219 0, /* tp_members */
220 0, /* tp_getset */
221 0, /* tp_base */
222 0, /* tp_dict */
223 0, /* tp_descr_get */
224 0, /* tp_descr_set */
225 0, /* tp_dictoffset */
226 0, /* tp_init */
227 0, /* tp_alloc */
228 groupby_new, /* tp_new */
229 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000230};
231
232
233/* _grouper object (internal) ************************************************/
234
235typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000236 PyObject_HEAD
237 PyObject *parent;
238 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000239} _grouperobject;
240
241static PyTypeObject _grouper_type;
242
243static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000244_grouper_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
245{
246 PyObject *parent, *tgtkey;
247
248 if (!PyArg_ParseTuple(args, "O!O", &groupby_type, &parent, &tgtkey))
249 return NULL;
250
251 return _grouper_create((groupbyobject*) parent, tgtkey);
252}
253
254static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000255_grouper_create(groupbyobject *parent, PyObject *tgtkey)
256{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000257 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000259 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
260 if (igo == NULL)
261 return NULL;
262 igo->parent = (PyObject *)parent;
263 Py_INCREF(parent);
264 igo->tgtkey = tgtkey;
265 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000267 PyObject_GC_Track(igo);
268 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000269}
270
271static void
272_grouper_dealloc(_grouperobject *igo)
273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000274 PyObject_GC_UnTrack(igo);
275 Py_DECREF(igo->parent);
276 Py_DECREF(igo->tgtkey);
277 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000278}
279
280static int
281_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
282{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000283 Py_VISIT(igo->parent);
284 Py_VISIT(igo->tgtkey);
285 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000286}
287
288static PyObject *
289_grouper_next(_grouperobject *igo)
290{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000291 groupbyobject *gbo = (groupbyobject *)igo->parent;
292 PyObject *newvalue, *newkey, *r;
293 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000295 if (gbo->currvalue == NULL) {
296 newvalue = PyIter_Next(gbo->it);
297 if (newvalue == NULL)
298 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000300 if (gbo->keyfunc == Py_None) {
301 newkey = newvalue;
302 Py_INCREF(newvalue);
303 } else {
304 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
305 newvalue, NULL);
306 if (newkey == NULL) {
307 Py_DECREF(newvalue);
308 return NULL;
309 }
310 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 assert(gbo->currkey == NULL);
313 gbo->currkey = newkey;
314 gbo->currvalue = newvalue;
315 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000317 assert(gbo->currkey != NULL);
318 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
319 if (rcmp <= 0)
320 /* got any error or current group is end */
321 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000322
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000323 r = gbo->currvalue;
324 gbo->currvalue = NULL;
325 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000326
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000327 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000328}
329
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000330static PyObject *
331_grouper_reduce(_grouperobject *lz)
332{
333 return Py_BuildValue("O(OO)", Py_TYPE(lz),
334 lz->parent, lz->tgtkey);
335}
336
337static PyMethodDef _grouper_methods[] = {
338 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
339 reduce_doc},
340 {NULL, NULL} /* sentinel */
341};
342
343
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000344static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000345 PyVarObject_HEAD_INIT(NULL, 0)
346 "itertools._grouper", /* tp_name */
347 sizeof(_grouperobject), /* tp_basicsize */
348 0, /* tp_itemsize */
349 /* methods */
350 (destructor)_grouper_dealloc, /* tp_dealloc */
351 0, /* tp_print */
352 0, /* tp_getattr */
353 0, /* tp_setattr */
354 0, /* tp_reserved */
355 0, /* tp_repr */
356 0, /* tp_as_number */
357 0, /* tp_as_sequence */
358 0, /* tp_as_mapping */
359 0, /* tp_hash */
360 0, /* tp_call */
361 0, /* tp_str */
362 PyObject_GenericGetAttr, /* tp_getattro */
363 0, /* tp_setattro */
364 0, /* tp_as_buffer */
365 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
366 0, /* tp_doc */
367 (traverseproc)_grouper_traverse,/* tp_traverse */
368 0, /* tp_clear */
369 0, /* tp_richcompare */
370 0, /* tp_weaklistoffset */
371 PyObject_SelfIter, /* tp_iter */
372 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000373 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000374 0, /* tp_members */
375 0, /* tp_getset */
376 0, /* tp_base */
377 0, /* tp_dict */
378 0, /* tp_descr_get */
379 0, /* tp_descr_set */
380 0, /* tp_dictoffset */
381 0, /* tp_init */
382 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000383 _grouper_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000384 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000385};
386
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000387
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000390
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391/* The teedataobject pre-allocates space for LINKCELLS number of objects.
392 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000393 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394 the other structure members including PyHEAD overhead. The larger the
395 value, the less memory overhead per object and the less time spent
396 allocating/deallocating new links. The smaller the number, the less
397 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400
401typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000402 PyObject_HEAD
403 PyObject *it;
404 int numread;
405 PyObject *nextlink;
406 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000407} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000408
409typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 PyObject_HEAD
411 teedataobject *dataobj;
412 int index;
413 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417
418static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000419teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000421 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000423 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
424 if (tdo == NULL)
425 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427 tdo->numread = 0;
428 tdo->nextlink = NULL;
429 Py_INCREF(it);
430 tdo->it = it;
431 PyObject_GC_Track(tdo);
432 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433}
434
435static PyObject *
436teedataobject_jumplink(teedataobject *tdo)
437{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000438 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000439 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000440 Py_XINCREF(tdo->nextlink);
441 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000442}
443
444static PyObject *
445teedataobject_getitem(teedataobject *tdo, int i)
446{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000449 assert(i < LINKCELLS);
450 if (i < tdo->numread)
451 value = tdo->values[i];
452 else {
453 /* this is the lead iterator, so fetch more data */
454 assert(i == tdo->numread);
455 value = PyIter_Next(tdo->it);
456 if (value == NULL)
457 return NULL;
458 tdo->numread++;
459 tdo->values[i] = value;
460 }
461 Py_INCREF(value);
462 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463}
464
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000465static int
466teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
467{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000468 int i;
469 Py_VISIT(tdo->it);
470 for (i = 0; i < tdo->numread; i++)
471 Py_VISIT(tdo->values[i]);
472 Py_VISIT(tdo->nextlink);
473 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000474}
475
476static int
477teedataobject_clear(teedataobject *tdo)
478{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000479 int i;
480 Py_CLEAR(tdo->it);
481 for (i=0 ; i<tdo->numread ; i++)
482 Py_CLEAR(tdo->values[i]);
483 Py_CLEAR(tdo->nextlink);
484 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000485}
486
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000487static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000488teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 PyObject_GC_UnTrack(tdo);
491 teedataobject_clear(tdo);
492 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000493}
494
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000495static PyObject *
496teedataobject_reduce(teedataobject *tdo)
497{
498 int i;
499 /* create a temporary list of already iterated values */
500 PyObject *values = PyList_New(tdo->numread);
501 if (!values)
502 return NULL;
503 for (i=0 ; i<tdo->numread ; i++) {
504 Py_INCREF(tdo->values[i]);
505 PyList_SET_ITEM(values, i, tdo->values[i]);
506 }
507 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
508 values,
509 tdo->nextlink ? tdo->nextlink : Py_None);
510}
511
512static PyTypeObject teedataobject_type;
513
514static PyObject *
515teedataobject_new(PyTypeObject *type, PyObject *args, PyObject *kw)
516{
517 teedataobject *tdo;
518 PyObject *it, *values, *next;
519 Py_ssize_t i, len;
520
521 assert(type == &teedataobject_type);
522 if (!PyArg_ParseTuple(args, "OO!O", &it, &PyList_Type, &values, &next))
523 return NULL;
524
525 tdo = (teedataobject *)teedataobject_newinternal(it);
526 if (!tdo)
527 return NULL;
528
529 len = PyList_GET_SIZE(values);
530 if (len > LINKCELLS)
531 goto err;
532 for (i=0; i<len; i++) {
533 tdo->values[i] = PyList_GET_ITEM(values, i);
534 Py_INCREF(tdo->values[i]);
535 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200536 /* len <= LINKCELLS < INT_MAX */
537 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000538
539 if (len == LINKCELLS) {
540 if (next != Py_None) {
541 if (Py_TYPE(next) != &teedataobject_type)
542 goto err;
543 assert(tdo->nextlink == NULL);
544 Py_INCREF(next);
545 tdo->nextlink = next;
546 }
547 } else {
548 if (next != Py_None)
549 goto err; /* shouldn't have a next if we are not full */
550 }
551 return (PyObject*)tdo;
552
553err:
554 Py_XDECREF(tdo);
555 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
556 return NULL;
557}
558
559static PyMethodDef teedataobject_methods[] = {
560 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
561 reduce_doc},
562 {NULL, NULL} /* sentinel */
563};
564
Raymond Hettingerad983e72003-11-12 14:32:26 +0000565PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000566
Raymond Hettingerad983e72003-11-12 14:32:26 +0000567static PyTypeObject teedataobject_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000568 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000569 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000570 sizeof(teedataobject), /* tp_basicsize */
571 0, /* tp_itemsize */
572 /* methods */
573 (destructor)teedataobject_dealloc, /* tp_dealloc */
574 0, /* tp_print */
575 0, /* tp_getattr */
576 0, /* tp_setattr */
577 0, /* tp_reserved */
578 0, /* tp_repr */
579 0, /* tp_as_number */
580 0, /* tp_as_sequence */
581 0, /* tp_as_mapping */
582 0, /* tp_hash */
583 0, /* tp_call */
584 0, /* tp_str */
585 PyObject_GenericGetAttr, /* tp_getattro */
586 0, /* tp_setattro */
587 0, /* tp_as_buffer */
588 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
589 teedataobject_doc, /* tp_doc */
590 (traverseproc)teedataobject_traverse, /* tp_traverse */
591 (inquiry)teedataobject_clear, /* tp_clear */
592 0, /* tp_richcompare */
593 0, /* tp_weaklistoffset */
594 0, /* tp_iter */
595 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000596 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000597 0, /* tp_members */
598 0, /* tp_getset */
599 0, /* tp_base */
600 0, /* tp_dict */
601 0, /* tp_descr_get */
602 0, /* tp_descr_set */
603 0, /* tp_dictoffset */
604 0, /* tp_init */
605 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000606 teedataobject_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000607 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000608};
609
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000610
611static PyTypeObject tee_type;
612
613static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000614tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000615{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000616 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000617
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000618 if (to->index >= LINKCELLS) {
619 link = teedataobject_jumplink(to->dataobj);
620 Py_DECREF(to->dataobj);
621 to->dataobj = (teedataobject *)link;
622 to->index = 0;
623 }
624 value = teedataobject_getitem(to->dataobj, to->index);
625 if (value == NULL)
626 return NULL;
627 to->index++;
628 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629}
630
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000631static int
632tee_traverse(teeobject *to, visitproc visit, void *arg)
633{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000634 Py_VISIT((PyObject *)to->dataobj);
635 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000636}
637
Raymond Hettingerad983e72003-11-12 14:32:26 +0000638static PyObject *
639tee_copy(teeobject *to)
640{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000641 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000642
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000643 newto = PyObject_GC_New(teeobject, &tee_type);
644 if (newto == NULL)
645 return NULL;
646 Py_INCREF(to->dataobj);
647 newto->dataobj = to->dataobj;
648 newto->index = to->index;
649 newto->weakreflist = NULL;
650 PyObject_GC_Track(newto);
651 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000652}
653
654PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
655
656static PyObject *
657tee_fromiterable(PyObject *iterable)
658{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000659 teeobject *to;
660 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000661
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000662 it = PyObject_GetIter(iterable);
663 if (it == NULL)
664 return NULL;
665 if (PyObject_TypeCheck(it, &tee_type)) {
666 to = (teeobject *)tee_copy((teeobject *)it);
667 goto done;
668 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000669
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000670 to = PyObject_GC_New(teeobject, &tee_type);
671 if (to == NULL)
672 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000673 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000674 if (!to->dataobj) {
675 PyObject_GC_Del(to);
676 to = NULL;
677 goto done;
678 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000679
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000680 to->index = 0;
681 to->weakreflist = NULL;
682 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000683done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000684 Py_XDECREF(it);
685 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000686}
687
688static PyObject *
689tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
690{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000691 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000692
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000693 if (!PyArg_UnpackTuple(args, "_tee", 1, 1, &iterable))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000694 return NULL;
695 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000696}
697
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000698static int
699tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000700{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000701 if (to->weakreflist != NULL)
702 PyObject_ClearWeakRefs((PyObject *) to);
703 Py_CLEAR(to->dataobj);
704 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000705}
706
707static void
708tee_dealloc(teeobject *to)
709{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000710 PyObject_GC_UnTrack(to);
711 tee_clear(to);
712 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000713}
714
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000715static PyObject *
716tee_reduce(teeobject *to)
717{
718 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
719}
720
721static PyObject *
722tee_setstate(teeobject *to, PyObject *state)
723{
724 teedataobject *tdo;
725 int index;
726 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
727 return NULL;
728 if (index < 0 || index > LINKCELLS) {
729 PyErr_SetString(PyExc_ValueError, "Index out of range");
730 return NULL;
731 }
732 Py_CLEAR(to->dataobj);
733 to->dataobj = tdo;
734 Py_INCREF(to->dataobj);
735 to->index = index;
736 Py_RETURN_NONE;
737}
738
Raymond Hettingerad983e72003-11-12 14:32:26 +0000739PyDoc_STRVAR(teeobject_doc,
740"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000741
Raymond Hettingerad983e72003-11-12 14:32:26 +0000742static PyMethodDef tee_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000743 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000744 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
745 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000746 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000747};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000748
749static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000750 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000751 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000752 sizeof(teeobject), /* tp_basicsize */
753 0, /* tp_itemsize */
754 /* methods */
755 (destructor)tee_dealloc, /* tp_dealloc */
756 0, /* tp_print */
757 0, /* tp_getattr */
758 0, /* tp_setattr */
759 0, /* tp_reserved */
760 0, /* tp_repr */
761 0, /* tp_as_number */
762 0, /* tp_as_sequence */
763 0, /* tp_as_mapping */
764 0, /* tp_hash */
765 0, /* tp_call */
766 0, /* tp_str */
767 0, /* tp_getattro */
768 0, /* tp_setattro */
769 0, /* tp_as_buffer */
770 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
771 teeobject_doc, /* tp_doc */
772 (traverseproc)tee_traverse, /* tp_traverse */
773 (inquiry)tee_clear, /* tp_clear */
774 0, /* tp_richcompare */
775 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
776 PyObject_SelfIter, /* tp_iter */
777 (iternextfunc)tee_next, /* tp_iternext */
778 tee_methods, /* tp_methods */
779 0, /* tp_members */
780 0, /* tp_getset */
781 0, /* tp_base */
782 0, /* tp_dict */
783 0, /* tp_descr_get */
784 0, /* tp_descr_set */
785 0, /* tp_dictoffset */
786 0, /* tp_init */
787 0, /* tp_alloc */
788 tee_new, /* tp_new */
789 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000790};
791
Raymond Hettingerad983e72003-11-12 14:32:26 +0000792static PyObject *
793tee(PyObject *self, PyObject *args)
794{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000795 Py_ssize_t i, n=2;
796 PyObject *it, *iterable, *copyable, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200797 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000798
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000799 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
800 return NULL;
801 if (n < 0) {
802 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
803 return NULL;
804 }
805 result = PyTuple_New(n);
806 if (result == NULL)
807 return NULL;
808 if (n == 0)
809 return result;
810 it = PyObject_GetIter(iterable);
811 if (it == NULL) {
812 Py_DECREF(result);
813 return NULL;
814 }
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200815 if (!_PyObject_HasAttrId(it, &PyId___copy__)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000816 copyable = tee_fromiterable(it);
817 Py_DECREF(it);
818 if (copyable == NULL) {
819 Py_DECREF(result);
820 return NULL;
821 }
822 } else
823 copyable = it;
824 PyTuple_SET_ITEM(result, 0, copyable);
825 for (i=1 ; i<n ; i++) {
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200826
827 copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000828 if (copyable == NULL) {
829 Py_DECREF(result);
830 return NULL;
831 }
832 PyTuple_SET_ITEM(result, i, copyable);
833 }
834 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000835}
836
837PyDoc_STRVAR(tee_doc,
838"tee(iterable, n=2) --> tuple of n independent iterators.");
839
840
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000841/* cycle object **********************************************************/
842
843typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000844 PyObject_HEAD
845 PyObject *it;
846 PyObject *saved;
847 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000848} cycleobject;
849
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000850static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000851
852static PyObject *
853cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
854{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000855 PyObject *it;
856 PyObject *iterable;
857 PyObject *saved;
858 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000859
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000860 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
861 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000862
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000863 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
864 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000865
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000866 /* Get iterator. */
867 it = PyObject_GetIter(iterable);
868 if (it == NULL)
869 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000870
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000871 saved = PyList_New(0);
872 if (saved == NULL) {
873 Py_DECREF(it);
874 return NULL;
875 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000877 /* create cycleobject structure */
878 lz = (cycleobject *)type->tp_alloc(type, 0);
879 if (lz == NULL) {
880 Py_DECREF(it);
881 Py_DECREF(saved);
882 return NULL;
883 }
884 lz->it = it;
885 lz->saved = saved;
886 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000887
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000888 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000889}
890
891static void
892cycle_dealloc(cycleobject *lz)
893{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000894 PyObject_GC_UnTrack(lz);
895 Py_XDECREF(lz->saved);
896 Py_XDECREF(lz->it);
897 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000898}
899
900static int
901cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
902{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000903 Py_VISIT(lz->it);
904 Py_VISIT(lz->saved);
905 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000906}
907
908static PyObject *
909cycle_next(cycleobject *lz)
910{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000911 PyObject *item;
912 PyObject *it;
913 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000914
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000915 while (1) {
916 item = PyIter_Next(lz->it);
917 if (item != NULL) {
918 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
919 Py_DECREF(item);
920 return NULL;
921 }
922 return item;
923 }
924 if (PyErr_Occurred()) {
925 if (PyErr_ExceptionMatches(PyExc_StopIteration))
926 PyErr_Clear();
927 else
928 return NULL;
929 }
930 if (PyList_Size(lz->saved) == 0)
931 return NULL;
932 it = PyObject_GetIter(lz->saved);
933 if (it == NULL)
934 return NULL;
935 tmp = lz->it;
936 lz->it = it;
937 lz->firstpass = 1;
938 Py_DECREF(tmp);
939 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000940}
941
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000942static PyObject *
943cycle_reduce(cycleobject *lz)
944{
945 /* Create a new cycle with the iterator tuple, then set
946 * the saved state on it.
947 */
948 return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
949 lz->it, lz->saved, lz->firstpass);
950 }
951
952static PyObject *
953cycle_setstate(cycleobject *lz, PyObject *state)
954{
955 PyObject *saved=NULL;
956 int firstpass;
957 if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
958 return NULL;
959 Py_CLEAR(lz->saved);
960 lz->saved = saved;
961 Py_XINCREF(lz->saved);
962 lz->firstpass = firstpass != 0;
963 Py_RETURN_NONE;
964}
965
966static PyMethodDef cycle_methods[] = {
967 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
968 reduce_doc},
969 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
970 setstate_doc},
971 {NULL, NULL} /* sentinel */
972};
973
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000974PyDoc_STRVAR(cycle_doc,
975"cycle(iterable) --> cycle object\n\
976\n\
977Return elements from the iterable until it is exhausted.\n\
978Then repeat the sequence indefinitely.");
979
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000980static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000981 PyVarObject_HEAD_INIT(NULL, 0)
982 "itertools.cycle", /* tp_name */
983 sizeof(cycleobject), /* tp_basicsize */
984 0, /* tp_itemsize */
985 /* methods */
986 (destructor)cycle_dealloc, /* tp_dealloc */
987 0, /* tp_print */
988 0, /* tp_getattr */
989 0, /* tp_setattr */
990 0, /* tp_reserved */
991 0, /* tp_repr */
992 0, /* tp_as_number */
993 0, /* tp_as_sequence */
994 0, /* tp_as_mapping */
995 0, /* tp_hash */
996 0, /* tp_call */
997 0, /* tp_str */
998 PyObject_GenericGetAttr, /* tp_getattro */
999 0, /* tp_setattro */
1000 0, /* tp_as_buffer */
1001 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1002 Py_TPFLAGS_BASETYPE, /* tp_flags */
1003 cycle_doc, /* tp_doc */
1004 (traverseproc)cycle_traverse, /* tp_traverse */
1005 0, /* tp_clear */
1006 0, /* tp_richcompare */
1007 0, /* tp_weaklistoffset */
1008 PyObject_SelfIter, /* tp_iter */
1009 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001010 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001011 0, /* tp_members */
1012 0, /* tp_getset */
1013 0, /* tp_base */
1014 0, /* tp_dict */
1015 0, /* tp_descr_get */
1016 0, /* tp_descr_set */
1017 0, /* tp_dictoffset */
1018 0, /* tp_init */
1019 0, /* tp_alloc */
1020 cycle_new, /* tp_new */
1021 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001022};
1023
1024
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025/* dropwhile object **********************************************************/
1026
1027typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001028 PyObject_HEAD
1029 PyObject *func;
1030 PyObject *it;
1031 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001032} dropwhileobject;
1033
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001034static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001035
1036static PyObject *
1037dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1038{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001039 PyObject *func, *seq;
1040 PyObject *it;
1041 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001042
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001043 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
1044 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001046 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
1047 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001048
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001049 /* Get iterator. */
1050 it = PyObject_GetIter(seq);
1051 if (it == NULL)
1052 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001054 /* create dropwhileobject structure */
1055 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1056 if (lz == NULL) {
1057 Py_DECREF(it);
1058 return NULL;
1059 }
1060 Py_INCREF(func);
1061 lz->func = func;
1062 lz->it = it;
1063 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001065 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066}
1067
1068static void
1069dropwhile_dealloc(dropwhileobject *lz)
1070{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001071 PyObject_GC_UnTrack(lz);
1072 Py_XDECREF(lz->func);
1073 Py_XDECREF(lz->it);
1074 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075}
1076
1077static int
1078dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1079{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001080 Py_VISIT(lz->it);
1081 Py_VISIT(lz->func);
1082 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083}
1084
1085static PyObject *
1086dropwhile_next(dropwhileobject *lz)
1087{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001088 PyObject *item, *good;
1089 PyObject *it = lz->it;
1090 long ok;
1091 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001092
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001093 iternext = *Py_TYPE(it)->tp_iternext;
1094 for (;;) {
1095 item = iternext(it);
1096 if (item == NULL)
1097 return NULL;
1098 if (lz->start == 1)
1099 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001101 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1102 if (good == NULL) {
1103 Py_DECREF(item);
1104 return NULL;
1105 }
1106 ok = PyObject_IsTrue(good);
1107 Py_DECREF(good);
1108 if (!ok) {
1109 lz->start = 1;
1110 return item;
1111 }
1112 Py_DECREF(item);
1113 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001114}
1115
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001116static PyObject *
1117dropwhile_reduce(dropwhileobject *lz)
1118{
1119 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1120 lz->func, lz->it, lz->start);
1121}
1122
1123static PyObject *
1124dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1125{
1126 int start = PyObject_IsTrue(state);
1127 if (start == -1)
1128 return NULL;
1129 lz->start = start;
1130 Py_RETURN_NONE;
1131}
1132
1133static PyMethodDef dropwhile_methods[] = {
1134 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1135 reduce_doc},
1136 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1137 setstate_doc},
1138 {NULL, NULL} /* sentinel */
1139};
1140
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001141PyDoc_STRVAR(dropwhile_doc,
1142"dropwhile(predicate, iterable) --> dropwhile object\n\
1143\n\
1144Drop items from the iterable while predicate(item) is true.\n\
1145Afterwards, return every element until the iterable is exhausted.");
1146
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001147static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001148 PyVarObject_HEAD_INIT(NULL, 0)
1149 "itertools.dropwhile", /* tp_name */
1150 sizeof(dropwhileobject), /* tp_basicsize */
1151 0, /* tp_itemsize */
1152 /* methods */
1153 (destructor)dropwhile_dealloc, /* tp_dealloc */
1154 0, /* tp_print */
1155 0, /* tp_getattr */
1156 0, /* tp_setattr */
1157 0, /* tp_reserved */
1158 0, /* tp_repr */
1159 0, /* tp_as_number */
1160 0, /* tp_as_sequence */
1161 0, /* tp_as_mapping */
1162 0, /* tp_hash */
1163 0, /* tp_call */
1164 0, /* tp_str */
1165 PyObject_GenericGetAttr, /* tp_getattro */
1166 0, /* tp_setattro */
1167 0, /* tp_as_buffer */
1168 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1169 Py_TPFLAGS_BASETYPE, /* tp_flags */
1170 dropwhile_doc, /* tp_doc */
1171 (traverseproc)dropwhile_traverse, /* tp_traverse */
1172 0, /* tp_clear */
1173 0, /* tp_richcompare */
1174 0, /* tp_weaklistoffset */
1175 PyObject_SelfIter, /* tp_iter */
1176 (iternextfunc)dropwhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001177 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001178 0, /* tp_members */
1179 0, /* tp_getset */
1180 0, /* tp_base */
1181 0, /* tp_dict */
1182 0, /* tp_descr_get */
1183 0, /* tp_descr_set */
1184 0, /* tp_dictoffset */
1185 0, /* tp_init */
1186 0, /* tp_alloc */
1187 dropwhile_new, /* tp_new */
1188 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001189};
1190
1191
1192/* takewhile object **********************************************************/
1193
1194typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001195 PyObject_HEAD
1196 PyObject *func;
1197 PyObject *it;
1198 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001199} takewhileobject;
1200
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001201static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202
1203static PyObject *
1204takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1205{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001206 PyObject *func, *seq;
1207 PyObject *it;
1208 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001209
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001210 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
1211 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001212
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001213 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
1214 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001216 /* Get iterator. */
1217 it = PyObject_GetIter(seq);
1218 if (it == NULL)
1219 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001221 /* create takewhileobject structure */
1222 lz = (takewhileobject *)type->tp_alloc(type, 0);
1223 if (lz == NULL) {
1224 Py_DECREF(it);
1225 return NULL;
1226 }
1227 Py_INCREF(func);
1228 lz->func = func;
1229 lz->it = it;
1230 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001232 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001233}
1234
1235static void
1236takewhile_dealloc(takewhileobject *lz)
1237{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001238 PyObject_GC_UnTrack(lz);
1239 Py_XDECREF(lz->func);
1240 Py_XDECREF(lz->it);
1241 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242}
1243
1244static int
1245takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1246{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001247 Py_VISIT(lz->it);
1248 Py_VISIT(lz->func);
1249 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001250}
1251
1252static PyObject *
1253takewhile_next(takewhileobject *lz)
1254{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001255 PyObject *item, *good;
1256 PyObject *it = lz->it;
1257 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001259 if (lz->stop == 1)
1260 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001262 item = (*Py_TYPE(it)->tp_iternext)(it);
1263 if (item == NULL)
1264 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001265
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1267 if (good == NULL) {
1268 Py_DECREF(item);
1269 return NULL;
1270 }
1271 ok = PyObject_IsTrue(good);
1272 Py_DECREF(good);
1273 if (ok)
1274 return item;
1275 Py_DECREF(item);
1276 lz->stop = 1;
1277 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001278}
1279
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001280static PyObject *
1281takewhile_reduce(takewhileobject *lz)
1282{
1283 return Py_BuildValue("O(OO)l", Py_TYPE(lz),
1284 lz->func, lz->it, lz->stop);
1285}
1286
1287static PyObject *
1288takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1289{
1290 int stop = PyObject_IsTrue(state);
1291 if (stop == -1)
1292 return NULL;
1293 lz->stop = stop;
1294 Py_RETURN_NONE;
1295}
1296
1297static PyMethodDef takewhile_reduce_methods[] = {
1298 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1299 reduce_doc},
1300 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1301 setstate_doc},
1302 {NULL, NULL} /* sentinel */
1303};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304PyDoc_STRVAR(takewhile_doc,
1305"takewhile(predicate, iterable) --> takewhile object\n\
1306\n\
1307Return successive entries from an iterable as long as the \n\
1308predicate evaluates to true for each entry.");
1309
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001310static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001311 PyVarObject_HEAD_INIT(NULL, 0)
1312 "itertools.takewhile", /* tp_name */
1313 sizeof(takewhileobject), /* tp_basicsize */
1314 0, /* tp_itemsize */
1315 /* methods */
1316 (destructor)takewhile_dealloc, /* tp_dealloc */
1317 0, /* tp_print */
1318 0, /* tp_getattr */
1319 0, /* tp_setattr */
1320 0, /* tp_reserved */
1321 0, /* tp_repr */
1322 0, /* tp_as_number */
1323 0, /* tp_as_sequence */
1324 0, /* tp_as_mapping */
1325 0, /* tp_hash */
1326 0, /* tp_call */
1327 0, /* tp_str */
1328 PyObject_GenericGetAttr, /* tp_getattro */
1329 0, /* tp_setattro */
1330 0, /* tp_as_buffer */
1331 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1332 Py_TPFLAGS_BASETYPE, /* tp_flags */
1333 takewhile_doc, /* tp_doc */
1334 (traverseproc)takewhile_traverse, /* tp_traverse */
1335 0, /* tp_clear */
1336 0, /* tp_richcompare */
1337 0, /* tp_weaklistoffset */
1338 PyObject_SelfIter, /* tp_iter */
1339 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001340 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001341 0, /* tp_members */
1342 0, /* tp_getset */
1343 0, /* tp_base */
1344 0, /* tp_dict */
1345 0, /* tp_descr_get */
1346 0, /* tp_descr_set */
1347 0, /* tp_dictoffset */
1348 0, /* tp_init */
1349 0, /* tp_alloc */
1350 takewhile_new, /* tp_new */
1351 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352};
1353
1354
1355/* islice object ************************************************************/
1356
1357typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001358 PyObject_HEAD
1359 PyObject *it;
1360 Py_ssize_t next;
1361 Py_ssize_t stop;
1362 Py_ssize_t step;
1363 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364} isliceobject;
1365
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001366static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367
1368static PyObject *
1369islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1370{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001371 PyObject *seq;
1372 Py_ssize_t start=0, stop=-1, step=1;
1373 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1374 Py_ssize_t numargs;
1375 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001377 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1378 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001379
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001380 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1381 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001382
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001383 numargs = PyTuple_Size(args);
1384 if (numargs == 2) {
1385 if (a1 != Py_None) {
1386 stop = PyLong_AsSsize_t(a1);
1387 if (stop == -1) {
1388 if (PyErr_Occurred())
1389 PyErr_Clear();
1390 PyErr_SetString(PyExc_ValueError,
1391 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1392 return NULL;
1393 }
1394 }
1395 } else {
1396 if (a1 != Py_None)
1397 start = PyLong_AsSsize_t(a1);
1398 if (start == -1 && PyErr_Occurred())
1399 PyErr_Clear();
1400 if (a2 != Py_None) {
1401 stop = PyLong_AsSsize_t(a2);
1402 if (stop == -1) {
1403 if (PyErr_Occurred())
1404 PyErr_Clear();
1405 PyErr_SetString(PyExc_ValueError,
1406 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1407 return NULL;
1408 }
1409 }
1410 }
1411 if (start<0 || stop<-1) {
1412 PyErr_SetString(PyExc_ValueError,
1413 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1414 return NULL;
1415 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001416
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001417 if (a3 != NULL) {
1418 if (a3 != Py_None)
1419 step = PyLong_AsSsize_t(a3);
1420 if (step == -1 && PyErr_Occurred())
1421 PyErr_Clear();
1422 }
1423 if (step<1) {
1424 PyErr_SetString(PyExc_ValueError,
1425 "Step for islice() must be a positive integer or None.");
1426 return NULL;
1427 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001428
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001429 /* Get iterator. */
1430 it = PyObject_GetIter(seq);
1431 if (it == NULL)
1432 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001433
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001434 /* create isliceobject structure */
1435 lz = (isliceobject *)type->tp_alloc(type, 0);
1436 if (lz == NULL) {
1437 Py_DECREF(it);
1438 return NULL;
1439 }
1440 lz->it = it;
1441 lz->next = start;
1442 lz->stop = stop;
1443 lz->step = step;
1444 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001446 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447}
1448
1449static void
1450islice_dealloc(isliceobject *lz)
1451{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001452 PyObject_GC_UnTrack(lz);
1453 Py_XDECREF(lz->it);
1454 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455}
1456
1457static int
1458islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1459{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001460 Py_VISIT(lz->it);
1461 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462}
1463
1464static PyObject *
1465islice_next(isliceobject *lz)
1466{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001467 PyObject *item;
1468 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001469 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001470 Py_ssize_t oldnext;
1471 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001473 iternext = *Py_TYPE(it)->tp_iternext;
1474 while (lz->cnt < lz->next) {
1475 item = iternext(it);
1476 if (item == NULL)
1477 return NULL;
1478 Py_DECREF(item);
1479 lz->cnt++;
1480 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001481 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 return NULL;
1483 item = iternext(it);
1484 if (item == NULL)
1485 return NULL;
1486 lz->cnt++;
1487 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001488 /* The (size_t) cast below avoids the danger of undefined
1489 behaviour from signed integer overflow. */
1490 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001491 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1492 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001493 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001494}
1495
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001496static PyObject *
1497islice_reduce(isliceobject *lz)
1498{
1499 /* When unpickled, generate a new object with the same bounds,
1500 * then 'setstate' with the next and count
1501 */
1502 PyObject *stop;
1503 if (lz->stop == -1) {
1504 stop = Py_None;
1505 Py_INCREF(stop);
1506 } else {
1507 stop = PyLong_FromSsize_t(lz->stop);
1508 if (stop == NULL)
1509 return NULL;
1510 }
1511 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1512 lz->it, lz->next, stop, lz->step,
1513 lz->cnt);
1514}
1515
1516static PyObject *
1517islice_setstate(isliceobject *lz, PyObject *state)
1518{
1519 Py_ssize_t cnt = PyLong_AsSsize_t(state);
1520 if (cnt == -1 && PyErr_Occurred())
1521 return NULL;
1522 lz->cnt = cnt;
1523 Py_RETURN_NONE;
1524}
1525
1526static PyMethodDef islice_methods[] = {
1527 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1528 reduce_doc},
1529 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1530 setstate_doc},
1531 {NULL, NULL} /* sentinel */
1532};
1533
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001534PyDoc_STRVAR(islice_doc,
1535"islice(iterable, [start,] stop [, step]) --> islice object\n\
1536\n\
1537Return an iterator whose next() method returns selected values from an\n\
1538iterable. If start is specified, will skip all preceding elements;\n\
1539otherwise, start defaults to zero. Step defaults to one. If\n\
1540specified as another value, step determines how many values are \n\
1541skipped between successive calls. Works like a slice() on a list\n\
1542but returns an iterator.");
1543
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001544static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001545 PyVarObject_HEAD_INIT(NULL, 0)
1546 "itertools.islice", /* tp_name */
1547 sizeof(isliceobject), /* tp_basicsize */
1548 0, /* tp_itemsize */
1549 /* methods */
1550 (destructor)islice_dealloc, /* tp_dealloc */
1551 0, /* tp_print */
1552 0, /* tp_getattr */
1553 0, /* tp_setattr */
1554 0, /* tp_reserved */
1555 0, /* tp_repr */
1556 0, /* tp_as_number */
1557 0, /* tp_as_sequence */
1558 0, /* tp_as_mapping */
1559 0, /* tp_hash */
1560 0, /* tp_call */
1561 0, /* tp_str */
1562 PyObject_GenericGetAttr, /* tp_getattro */
1563 0, /* tp_setattro */
1564 0, /* tp_as_buffer */
1565 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1566 Py_TPFLAGS_BASETYPE, /* tp_flags */
1567 islice_doc, /* tp_doc */
1568 (traverseproc)islice_traverse, /* tp_traverse */
1569 0, /* tp_clear */
1570 0, /* tp_richcompare */
1571 0, /* tp_weaklistoffset */
1572 PyObject_SelfIter, /* tp_iter */
1573 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001574 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001575 0, /* tp_members */
1576 0, /* tp_getset */
1577 0, /* tp_base */
1578 0, /* tp_dict */
1579 0, /* tp_descr_get */
1580 0, /* tp_descr_set */
1581 0, /* tp_dictoffset */
1582 0, /* tp_init */
1583 0, /* tp_alloc */
1584 islice_new, /* tp_new */
1585 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001586};
1587
1588
1589/* starmap object ************************************************************/
1590
1591typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001592 PyObject_HEAD
1593 PyObject *func;
1594 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595} starmapobject;
1596
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001597static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001598
1599static PyObject *
1600starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1601{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001602 PyObject *func, *seq;
1603 PyObject *it;
1604 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001605
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001606 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1607 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001608
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001609 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1610 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001612 /* Get iterator. */
1613 it = PyObject_GetIter(seq);
1614 if (it == NULL)
1615 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001617 /* create starmapobject structure */
1618 lz = (starmapobject *)type->tp_alloc(type, 0);
1619 if (lz == NULL) {
1620 Py_DECREF(it);
1621 return NULL;
1622 }
1623 Py_INCREF(func);
1624 lz->func = func;
1625 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001626
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001627 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001628}
1629
1630static void
1631starmap_dealloc(starmapobject *lz)
1632{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001633 PyObject_GC_UnTrack(lz);
1634 Py_XDECREF(lz->func);
1635 Py_XDECREF(lz->it);
1636 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637}
1638
1639static int
1640starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1641{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001642 Py_VISIT(lz->it);
1643 Py_VISIT(lz->func);
1644 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645}
1646
1647static PyObject *
1648starmap_next(starmapobject *lz)
1649{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001650 PyObject *args;
1651 PyObject *result;
1652 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001653
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001654 args = (*Py_TYPE(it)->tp_iternext)(it);
1655 if (args == NULL)
1656 return NULL;
1657 if (!PyTuple_CheckExact(args)) {
1658 PyObject *newargs = PySequence_Tuple(args);
1659 Py_DECREF(args);
1660 if (newargs == NULL)
1661 return NULL;
1662 args = newargs;
1663 }
1664 result = PyObject_Call(lz->func, args, NULL);
1665 Py_DECREF(args);
1666 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001667}
1668
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001669static PyObject *
1670starmap_reduce(starmapobject *lz)
1671{
1672 /* Just pickle the iterator */
1673 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1674}
1675
1676static PyMethodDef starmap_methods[] = {
1677 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1678 reduce_doc},
1679 {NULL, NULL} /* sentinel */
1680};
1681
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682PyDoc_STRVAR(starmap_doc,
1683"starmap(function, sequence) --> starmap object\n\
1684\n\
1685Return an iterator whose values are returned from the function evaluated\n\
1686with a argument tuple taken from the given sequence.");
1687
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001688static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001689 PyVarObject_HEAD_INIT(NULL, 0)
1690 "itertools.starmap", /* tp_name */
1691 sizeof(starmapobject), /* tp_basicsize */
1692 0, /* tp_itemsize */
1693 /* methods */
1694 (destructor)starmap_dealloc, /* tp_dealloc */
1695 0, /* tp_print */
1696 0, /* tp_getattr */
1697 0, /* tp_setattr */
1698 0, /* tp_reserved */
1699 0, /* tp_repr */
1700 0, /* tp_as_number */
1701 0, /* tp_as_sequence */
1702 0, /* tp_as_mapping */
1703 0, /* tp_hash */
1704 0, /* tp_call */
1705 0, /* tp_str */
1706 PyObject_GenericGetAttr, /* tp_getattro */
1707 0, /* tp_setattro */
1708 0, /* tp_as_buffer */
1709 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1710 Py_TPFLAGS_BASETYPE, /* tp_flags */
1711 starmap_doc, /* tp_doc */
1712 (traverseproc)starmap_traverse, /* tp_traverse */
1713 0, /* tp_clear */
1714 0, /* tp_richcompare */
1715 0, /* tp_weaklistoffset */
1716 PyObject_SelfIter, /* tp_iter */
1717 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001718 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001719 0, /* tp_members */
1720 0, /* tp_getset */
1721 0, /* tp_base */
1722 0, /* tp_dict */
1723 0, /* tp_descr_get */
1724 0, /* tp_descr_set */
1725 0, /* tp_dictoffset */
1726 0, /* tp_init */
1727 0, /* tp_alloc */
1728 starmap_new, /* tp_new */
1729 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001730};
1731
1732
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001733/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001734
1735typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001736 PyObject_HEAD
1737 PyObject *source; /* Iterator over input iterables */
1738 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001739} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001740
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001741static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001743static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001744chain_new_internal(PyTypeObject *type, PyObject *source)
1745{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001746 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001747
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001748 lz = (chainobject *)type->tp_alloc(type, 0);
1749 if (lz == NULL) {
1750 Py_DECREF(source);
1751 return NULL;
1752 }
1753
1754 lz->source = source;
1755 lz->active = NULL;
1756 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001757}
1758
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001760chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001761{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001762 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001763
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001764 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1765 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001766
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001767 source = PyObject_GetIter(args);
1768 if (source == NULL)
1769 return NULL;
1770
1771 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001772}
1773
1774static PyObject *
1775chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1776{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001777 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001778
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001779 source = PyObject_GetIter(arg);
1780 if (source == NULL)
1781 return NULL;
1782
1783 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784}
1785
1786static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001787chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001788{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001789 PyObject_GC_UnTrack(lz);
1790 Py_XDECREF(lz->active);
1791 Py_XDECREF(lz->source);
1792 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001793}
1794
Raymond Hettinger2012f172003-02-07 05:32:58 +00001795static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001796chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001797{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001798 Py_VISIT(lz->source);
1799 Py_VISIT(lz->active);
1800 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001801}
1802
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001804chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001806 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001807
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001808 if (lz->source == NULL)
1809 return NULL; /* already stopped */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001810
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001811 if (lz->active == NULL) {
1812 PyObject *iterable = PyIter_Next(lz->source);
1813 if (iterable == NULL) {
1814 Py_CLEAR(lz->source);
1815 return NULL; /* no more input sources */
1816 }
1817 lz->active = PyObject_GetIter(iterable);
1818 Py_DECREF(iterable);
1819 if (lz->active == NULL) {
1820 Py_CLEAR(lz->source);
1821 return NULL; /* input not iterable */
1822 }
1823 }
1824 item = PyIter_Next(lz->active);
1825 if (item != NULL)
1826 return item;
1827 if (PyErr_Occurred()) {
1828 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1829 PyErr_Clear();
1830 else
1831 return NULL; /* input raised an exception */
1832 }
1833 Py_CLEAR(lz->active);
1834 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001835}
1836
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001837static PyObject *
1838chain_reduce(chainobject *lz)
1839{
1840 if (lz->source) {
1841 /* we can't pickle function objects (itertools.from_iterable) so
1842 * we must use setstate to replace the iterable. One day we
1843 * will fix pickling of functions
1844 */
1845 if (lz->active) {
1846 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1847 } else {
1848 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1849 }
1850 } else {
1851 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1852 }
1853 return NULL;
1854}
1855
1856static PyObject *
1857chain_setstate(chainobject *lz, PyObject *state)
1858{
1859 PyObject *source, *active=NULL;
1860 if (! PyArg_ParseTuple(state, "O|O", &source, &active))
1861 return NULL;
1862
1863 Py_CLEAR(lz->source);
1864 lz->source = source;
1865 Py_INCREF(lz->source);
1866 Py_CLEAR(lz->active);
1867 lz->active = active;
1868 Py_XINCREF(lz->active);
1869 Py_RETURN_NONE;
1870}
1871
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001872PyDoc_STRVAR(chain_doc,
1873"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001874\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00001875Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001876first iterable until it is exhausted, then elements from the next\n\
1877iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001878
Christian Heimesf16baeb2008-02-29 14:57:44 +00001879PyDoc_STRVAR(chain_from_iterable_doc,
1880"chain.from_iterable(iterable) --> chain object\n\
1881\n\
1882Alternate chain() contructor taking a single iterable argument\n\
1883that evaluates lazily.");
1884
1885static PyMethodDef chain_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001886 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1887 chain_from_iterable_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001888 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
1889 reduce_doc},
1890 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
1891 setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001892 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00001893};
1894
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001895static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001896 PyVarObject_HEAD_INIT(NULL, 0)
1897 "itertools.chain", /* tp_name */
1898 sizeof(chainobject), /* tp_basicsize */
1899 0, /* tp_itemsize */
1900 /* methods */
1901 (destructor)chain_dealloc, /* tp_dealloc */
1902 0, /* tp_print */
1903 0, /* tp_getattr */
1904 0, /* tp_setattr */
1905 0, /* tp_reserved */
1906 0, /* tp_repr */
1907 0, /* tp_as_number */
1908 0, /* tp_as_sequence */
1909 0, /* tp_as_mapping */
1910 0, /* tp_hash */
1911 0, /* tp_call */
1912 0, /* tp_str */
1913 PyObject_GenericGetAttr, /* tp_getattro */
1914 0, /* tp_setattro */
1915 0, /* tp_as_buffer */
1916 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1917 Py_TPFLAGS_BASETYPE, /* tp_flags */
1918 chain_doc, /* tp_doc */
1919 (traverseproc)chain_traverse, /* tp_traverse */
1920 0, /* tp_clear */
1921 0, /* tp_richcompare */
1922 0, /* tp_weaklistoffset */
1923 PyObject_SelfIter, /* tp_iter */
1924 (iternextfunc)chain_next, /* tp_iternext */
1925 chain_methods, /* tp_methods */
1926 0, /* tp_members */
1927 0, /* tp_getset */
1928 0, /* tp_base */
1929 0, /* tp_dict */
1930 0, /* tp_descr_get */
1931 0, /* tp_descr_set */
1932 0, /* tp_dictoffset */
1933 0, /* tp_init */
1934 0, /* tp_alloc */
1935 chain_new, /* tp_new */
1936 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001937};
1938
1939
Christian Heimesc3f30c42008-02-22 16:37:40 +00001940/* product object ************************************************************/
1941
1942typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001943 PyObject_HEAD
1944 PyObject *pools; /* tuple of pool tuples */
1945 Py_ssize_t *indices; /* one index per pool */
1946 PyObject *result; /* most recently returned result tuple */
1947 int stopped; /* set to 1 when the product iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00001948} productobject;
1949
1950static PyTypeObject product_type;
1951
1952static PyObject *
1953product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1954{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001955 productobject *lz;
1956 Py_ssize_t nargs, npools, repeat=1;
1957 PyObject *pools = NULL;
1958 Py_ssize_t *indices = NULL;
1959 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001960
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001961 if (kwds != NULL) {
1962 char *kwlist[] = {"repeat", 0};
1963 PyObject *tmpargs = PyTuple_New(0);
1964 if (tmpargs == NULL)
1965 return NULL;
1966 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1967 Py_DECREF(tmpargs);
1968 return NULL;
1969 }
1970 Py_DECREF(tmpargs);
1971 if (repeat < 0) {
1972 PyErr_SetString(PyExc_ValueError,
1973 "repeat argument cannot be negative");
1974 return NULL;
1975 }
1976 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001977
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001978 assert(PyTuple_Check(args));
1979 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1980 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001981
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001982 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1983 if (indices == NULL) {
1984 PyErr_NoMemory();
1985 goto error;
1986 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00001987
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001988 pools = PyTuple_New(npools);
1989 if (pools == NULL)
1990 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00001991
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001992 for (i=0; i < nargs ; ++i) {
1993 PyObject *item = PyTuple_GET_ITEM(args, i);
1994 PyObject *pool = PySequence_Tuple(item);
1995 if (pool == NULL)
1996 goto error;
1997 PyTuple_SET_ITEM(pools, i, pool);
1998 indices[i] = 0;
1999 }
2000 for ( ; i < npools; ++i) {
2001 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2002 Py_INCREF(pool);
2003 PyTuple_SET_ITEM(pools, i, pool);
2004 indices[i] = 0;
2005 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002007 /* create productobject structure */
2008 lz = (productobject *)type->tp_alloc(type, 0);
2009 if (lz == NULL)
2010 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002011
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002012 lz->pools = pools;
2013 lz->indices = indices;
2014 lz->result = NULL;
2015 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002016
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002017 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002018
2019error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002020 if (indices != NULL)
2021 PyMem_Free(indices);
2022 Py_XDECREF(pools);
2023 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002024}
2025
2026static void
2027product_dealloc(productobject *lz)
2028{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002029 PyObject_GC_UnTrack(lz);
2030 Py_XDECREF(lz->pools);
2031 Py_XDECREF(lz->result);
2032 if (lz->indices != NULL)
2033 PyMem_Free(lz->indices);
2034 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002035}
2036
2037static int
2038product_traverse(productobject *lz, visitproc visit, void *arg)
2039{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002040 Py_VISIT(lz->pools);
2041 Py_VISIT(lz->result);
2042 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002043}
2044
2045static PyObject *
2046product_next(productobject *lz)
2047{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002048 PyObject *pool;
2049 PyObject *elem;
2050 PyObject *oldelem;
2051 PyObject *pools = lz->pools;
2052 PyObject *result = lz->result;
2053 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2054 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002055
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002056 if (lz->stopped)
2057 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002058
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002059 if (result == NULL) {
2060 /* On the first pass, return an initial tuple filled with the
2061 first element from each pool. */
2062 result = PyTuple_New(npools);
2063 if (result == NULL)
2064 goto empty;
2065 lz->result = result;
2066 for (i=0; i < npools; i++) {
2067 pool = PyTuple_GET_ITEM(pools, i);
2068 if (PyTuple_GET_SIZE(pool) == 0)
2069 goto empty;
2070 elem = PyTuple_GET_ITEM(pool, 0);
2071 Py_INCREF(elem);
2072 PyTuple_SET_ITEM(result, i, elem);
2073 }
2074 } else {
2075 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002077 /* Copy the previous result tuple or re-use it if available */
2078 if (Py_REFCNT(result) > 1) {
2079 PyObject *old_result = result;
2080 result = PyTuple_New(npools);
2081 if (result == NULL)
2082 goto empty;
2083 lz->result = result;
2084 for (i=0; i < npools; i++) {
2085 elem = PyTuple_GET_ITEM(old_result, i);
2086 Py_INCREF(elem);
2087 PyTuple_SET_ITEM(result, i, elem);
2088 }
2089 Py_DECREF(old_result);
2090 }
2091 /* Now, we've got the only copy so we can update it in-place */
2092 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002093
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002094 /* Update the pool indices right-to-left. Only advance to the
2095 next pool when the previous one rolls-over */
2096 for (i=npools-1 ; i >= 0 ; i--) {
2097 pool = PyTuple_GET_ITEM(pools, i);
2098 indices[i]++;
2099 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2100 /* Roll-over and advance to next pool */
2101 indices[i] = 0;
2102 elem = PyTuple_GET_ITEM(pool, 0);
2103 Py_INCREF(elem);
2104 oldelem = PyTuple_GET_ITEM(result, i);
2105 PyTuple_SET_ITEM(result, i, elem);
2106 Py_DECREF(oldelem);
2107 } else {
2108 /* No rollover. Just increment and stop here. */
2109 elem = PyTuple_GET_ITEM(pool, indices[i]);
2110 Py_INCREF(elem);
2111 oldelem = PyTuple_GET_ITEM(result, i);
2112 PyTuple_SET_ITEM(result, i, elem);
2113 Py_DECREF(oldelem);
2114 break;
2115 }
2116 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002117
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002118 /* If i is negative, then the indices have all rolled-over
2119 and we're done. */
2120 if (i < 0)
2121 goto empty;
2122 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002123
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002124 Py_INCREF(result);
2125 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002126
2127empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002128 lz->stopped = 1;
2129 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002130}
2131
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002132static PyObject *
2133product_reduce(productobject *lz)
2134{
2135 if (lz->stopped) {
2136 return Py_BuildValue("O(())", Py_TYPE(lz));
2137 } else if (lz->result == NULL) {
2138 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2139 } else {
2140 PyObject *indices;
2141 Py_ssize_t n, i;
2142
2143 /* we must pickle the indices use them for setstate, and
2144 * additionally indicate that the iterator has started
2145 */
2146 n = PyTuple_GET_SIZE(lz->pools);
2147 indices = PyTuple_New(n);
2148 if (indices == NULL)
2149 return NULL;
2150 for (i=0; i<n; i++){
2151 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2152 if (!index) {
2153 Py_DECREF(indices);
2154 return NULL;
2155 }
2156 PyTuple_SET_ITEM(indices, i, index);
2157 }
2158 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2159 }
2160}
2161
2162static PyObject *
2163product_setstate(productobject *lz, PyObject *state)
2164{
2165 PyObject *result;
2166 Py_ssize_t n, i;
2167
2168 n = PyTuple_GET_SIZE(lz->pools);
2169 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2170 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2171 return NULL;
2172 }
2173 for (i=0; i<n; i++)
2174 {
2175 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2176 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2177 if (index < 0 && PyErr_Occurred())
2178 return NULL; /* not an integer */
2179 /* clamp the index */
2180 if (index < 0)
2181 index = 0;
2182 else if (index > n-1)
2183 index = n-1;
2184 lz->indices[i] = index;
2185 }
2186
2187 result = PyTuple_New(n);
2188 if (!result)
2189 return NULL;
2190 for (i=0; i<n; i++) {
2191 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2192 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2193 Py_INCREF(element);
2194 PyTuple_SET_ITEM(result, i, element);
2195 }
2196 Py_CLEAR(lz->result);
2197 lz->result = result;
2198 Py_RETURN_NONE;
2199}
2200
2201static PyMethodDef product_methods[] = {
2202 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2203 reduce_doc},
2204 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2205 setstate_doc},
2206 {NULL, NULL} /* sentinel */
2207};
2208
Christian Heimesc3f30c42008-02-22 16:37:40 +00002209PyDoc_STRVAR(product_doc,
2210"product(*iterables) --> product object\n\
2211\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002212Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002213For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2214The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2215cycle in a manner similar to an odometer (with the rightmost element changing\n\
2216on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002217To compute the product of an iterable with itself, specify the number\n\
2218of repetitions with the optional repeat keyword argument. For example,\n\
2219product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002220product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2221product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2222
2223static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002224 PyVarObject_HEAD_INIT(NULL, 0)
2225 "itertools.product", /* tp_name */
2226 sizeof(productobject), /* tp_basicsize */
2227 0, /* tp_itemsize */
2228 /* methods */
2229 (destructor)product_dealloc, /* tp_dealloc */
2230 0, /* tp_print */
2231 0, /* tp_getattr */
2232 0, /* tp_setattr */
2233 0, /* tp_reserved */
2234 0, /* tp_repr */
2235 0, /* tp_as_number */
2236 0, /* tp_as_sequence */
2237 0, /* tp_as_mapping */
2238 0, /* tp_hash */
2239 0, /* tp_call */
2240 0, /* tp_str */
2241 PyObject_GenericGetAttr, /* tp_getattro */
2242 0, /* tp_setattro */
2243 0, /* tp_as_buffer */
2244 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2245 Py_TPFLAGS_BASETYPE, /* tp_flags */
2246 product_doc, /* tp_doc */
2247 (traverseproc)product_traverse, /* tp_traverse */
2248 0, /* tp_clear */
2249 0, /* tp_richcompare */
2250 0, /* tp_weaklistoffset */
2251 PyObject_SelfIter, /* tp_iter */
2252 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002253 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002254 0, /* tp_members */
2255 0, /* tp_getset */
2256 0, /* tp_base */
2257 0, /* tp_dict */
2258 0, /* tp_descr_get */
2259 0, /* tp_descr_set */
2260 0, /* tp_dictoffset */
2261 0, /* tp_init */
2262 0, /* tp_alloc */
2263 product_new, /* tp_new */
2264 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002265};
2266
2267
Christian Heimes380f7f22008-02-28 11:19:05 +00002268/* combinations object ************************************************************/
2269
2270typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002271 PyObject_HEAD
2272 PyObject *pool; /* input converted to a tuple */
2273 Py_ssize_t *indices; /* one index per result element */
2274 PyObject *result; /* most recently returned result tuple */
2275 Py_ssize_t r; /* size of result tuple */
2276 int stopped; /* set to 1 when the combinations iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002277} combinationsobject;
2278
2279static PyTypeObject combinations_type;
2280
2281static PyObject *
2282combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2283{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002284 combinationsobject *co;
2285 Py_ssize_t n;
2286 Py_ssize_t r;
2287 PyObject *pool = NULL;
2288 PyObject *iterable = NULL;
2289 Py_ssize_t *indices = NULL;
2290 Py_ssize_t i;
2291 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimes380f7f22008-02-28 11:19:05 +00002292
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002293 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2294 &iterable, &r))
2295 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002296
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002297 pool = PySequence_Tuple(iterable);
2298 if (pool == NULL)
2299 goto error;
2300 n = PyTuple_GET_SIZE(pool);
2301 if (r < 0) {
2302 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2303 goto error;
2304 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002305
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002306 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2307 if (indices == NULL) {
2308 PyErr_NoMemory();
2309 goto error;
2310 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002312 for (i=0 ; i<r ; i++)
2313 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002314
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002315 /* create combinationsobject structure */
2316 co = (combinationsobject *)type->tp_alloc(type, 0);
2317 if (co == NULL)
2318 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002319
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002320 co->pool = pool;
2321 co->indices = indices;
2322 co->result = NULL;
2323 co->r = r;
2324 co->stopped = r > n ? 1 : 0;
2325
2326 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002327
2328error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002329 if (indices != NULL)
2330 PyMem_Free(indices);
2331 Py_XDECREF(pool);
2332 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002333}
2334
2335static void
2336combinations_dealloc(combinationsobject *co)
2337{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002338 PyObject_GC_UnTrack(co);
2339 Py_XDECREF(co->pool);
2340 Py_XDECREF(co->result);
2341 if (co->indices != NULL)
2342 PyMem_Free(co->indices);
2343 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002344}
2345
2346static int
2347combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2348{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002349 Py_VISIT(co->pool);
2350 Py_VISIT(co->result);
2351 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002352}
2353
2354static PyObject *
2355combinations_next(combinationsobject *co)
2356{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002357 PyObject *elem;
2358 PyObject *oldelem;
2359 PyObject *pool = co->pool;
2360 Py_ssize_t *indices = co->indices;
2361 PyObject *result = co->result;
2362 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2363 Py_ssize_t r = co->r;
2364 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002365
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002366 if (co->stopped)
2367 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002368
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002369 if (result == NULL) {
2370 /* On the first pass, initialize result tuple using the indices */
2371 result = PyTuple_New(r);
2372 if (result == NULL)
2373 goto empty;
2374 co->result = result;
2375 for (i=0; i<r ; i++) {
2376 index = indices[i];
2377 elem = PyTuple_GET_ITEM(pool, index);
2378 Py_INCREF(elem);
2379 PyTuple_SET_ITEM(result, i, elem);
2380 }
2381 } else {
2382 /* Copy the previous result tuple or re-use it if available */
2383 if (Py_REFCNT(result) > 1) {
2384 PyObject *old_result = result;
2385 result = PyTuple_New(r);
2386 if (result == NULL)
2387 goto empty;
2388 co->result = result;
2389 for (i=0; i<r ; i++) {
2390 elem = PyTuple_GET_ITEM(old_result, i);
2391 Py_INCREF(elem);
2392 PyTuple_SET_ITEM(result, i, elem);
2393 }
2394 Py_DECREF(old_result);
2395 }
2396 /* Now, we've got the only copy so we can update it in-place
2397 * CPython's empty tuple is a singleton and cached in
2398 * PyTuple's freelist.
2399 */
2400 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002401
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002402 /* Scan indices right-to-left until finding one that is not
2403 at its maximum (i + n - r). */
2404 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2405 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002406
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002407 /* If i is negative, then the indices are all at
2408 their maximum value and we're done. */
2409 if (i < 0)
2410 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002411
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002412 /* Increment the current index which we know is not at its
2413 maximum. Then move back to the right setting each index
2414 to its lowest possible value (one higher than the index
2415 to its left -- this maintains the sort order invariant). */
2416 indices[i]++;
2417 for (j=i+1 ; j<r ; j++)
2418 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002419
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002420 /* Update the result tuple for the new indices
2421 starting with i, the leftmost index that changed */
2422 for ( ; i<r ; i++) {
2423 index = indices[i];
2424 elem = PyTuple_GET_ITEM(pool, index);
2425 Py_INCREF(elem);
2426 oldelem = PyTuple_GET_ITEM(result, i);
2427 PyTuple_SET_ITEM(result, i, elem);
2428 Py_DECREF(oldelem);
2429 }
2430 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002431
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002432 Py_INCREF(result);
2433 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002434
2435empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002436 co->stopped = 1;
2437 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002438}
2439
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002440static PyObject *
2441combinations_reduce(combinationsobject *lz)
2442{
2443 if (lz->result == NULL) {
2444 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2445 } else if (lz->stopped) {
2446 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2447 } else {
2448 PyObject *indices;
2449 Py_ssize_t i;
2450
2451 /* we must pickle the indices and use them for setstate */
2452 indices = PyTuple_New(lz->r);
2453 if (!indices)
2454 return NULL;
2455 for (i=0; i<lz->r; i++)
2456 {
2457 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2458 if (!index) {
2459 Py_DECREF(indices);
2460 return NULL;
2461 }
2462 PyTuple_SET_ITEM(indices, i, index);
2463 }
2464
2465 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2466 }
2467}
2468
2469static PyObject *
2470combinations_setstate(combinationsobject *lz, PyObject *state)
2471{
2472 PyObject *result;
2473 Py_ssize_t i;
2474 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2475
2476 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2477 {
2478 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2479 return NULL;
2480 }
2481
2482 for (i=0; i<lz->r; i++)
2483 {
2484 Py_ssize_t max;
2485 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2486 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2487 if (index == -1 && PyErr_Occurred())
2488 return NULL; /* not an integer */
2489 max = i + n - lz->r;
2490 /* clamp the index (beware of negative max) */
2491 if (index > max)
2492 index = max;
2493 if (index < 0)
2494 index = 0;
2495 lz->indices[i] = index;
2496 }
2497
2498 result = PyTuple_New(lz->r);
2499 if (result == NULL)
2500 return NULL;
2501 for (i=0; i<lz->r; i++) {
2502 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2503 Py_INCREF(element);
2504 PyTuple_SET_ITEM(result, i, element);
2505 }
2506
2507 Py_CLEAR(lz->result);
2508 lz->result = result;
2509 Py_RETURN_NONE;
2510}
2511
2512static PyMethodDef combinations_methods[] = {
2513 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2514 reduce_doc},
2515 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2516 setstate_doc},
2517 {NULL, NULL} /* sentinel */
2518};
2519
Christian Heimes380f7f22008-02-28 11:19:05 +00002520PyDoc_STRVAR(combinations_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002521"combinations(iterable, r) --> combinations object\n\
Christian Heimes380f7f22008-02-28 11:19:05 +00002522\n\
2523Return successive r-length combinations of elements in the iterable.\n\n\
2524combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2525
2526static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002527 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002528 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002529 sizeof(combinationsobject), /* tp_basicsize */
2530 0, /* tp_itemsize */
2531 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002532 (destructor)combinations_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002533 0, /* tp_print */
2534 0, /* tp_getattr */
2535 0, /* tp_setattr */
2536 0, /* tp_reserved */
2537 0, /* tp_repr */
2538 0, /* tp_as_number */
2539 0, /* tp_as_sequence */
2540 0, /* tp_as_mapping */
2541 0, /* tp_hash */
2542 0, /* tp_call */
2543 0, /* tp_str */
2544 PyObject_GenericGetAttr, /* tp_getattro */
2545 0, /* tp_setattro */
2546 0, /* tp_as_buffer */
2547 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2548 Py_TPFLAGS_BASETYPE, /* tp_flags */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002549 combinations_doc, /* tp_doc */
2550 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002551 0, /* tp_clear */
2552 0, /* tp_richcompare */
2553 0, /* tp_weaklistoffset */
2554 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002555 (iternextfunc)combinations_next, /* tp_iternext */
2556 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002557 0, /* tp_members */
2558 0, /* tp_getset */
2559 0, /* tp_base */
2560 0, /* tp_dict */
2561 0, /* tp_descr_get */
2562 0, /* tp_descr_set */
2563 0, /* tp_dictoffset */
2564 0, /* tp_init */
2565 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002566 combinations_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002567 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002568};
2569
2570
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002571/* combinations with replacement object *******************************************/
2572
2573/* Equivalent to:
2574
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002575 def combinations_with_replacement(iterable, r):
2576 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2577 # number items returned: (n+r-1)! / r! / (n-1)!
2578 pool = tuple(iterable)
2579 n = len(pool)
2580 indices = [0] * r
2581 yield tuple(pool[i] for i in indices)
2582 while 1:
2583 for i in reversed(range(r)):
2584 if indices[i] != n - 1:
2585 break
2586 else:
2587 return
2588 indices[i:] = [indices[i] + 1] * (r - i)
2589 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002590
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002591 def combinations_with_replacement2(iterable, r):
2592 'Alternate version that filters from product()'
2593 pool = tuple(iterable)
2594 n = len(pool)
2595 for indices in product(range(n), repeat=r):
2596 if sorted(indices) == list(indices):
2597 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002598*/
2599typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002600 PyObject_HEAD
2601 PyObject *pool; /* input converted to a tuple */
2602 Py_ssize_t *indices; /* one index per result element */
2603 PyObject *result; /* most recently returned result tuple */
2604 Py_ssize_t r; /* size of result tuple */
2605 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002606} cwrobject;
2607
2608static PyTypeObject cwr_type;
2609
2610static PyObject *
2611cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2612{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002613 cwrobject *co;
2614 Py_ssize_t n;
2615 Py_ssize_t r;
2616 PyObject *pool = NULL;
2617 PyObject *iterable = NULL;
2618 Py_ssize_t *indices = NULL;
2619 Py_ssize_t i;
2620 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002621
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002622 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2623 &iterable, &r))
2624 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002625
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002626 pool = PySequence_Tuple(iterable);
2627 if (pool == NULL)
2628 goto error;
2629 n = PyTuple_GET_SIZE(pool);
2630 if (r < 0) {
2631 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2632 goto error;
2633 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002634
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002635 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2636 if (indices == NULL) {
2637 PyErr_NoMemory();
2638 goto error;
2639 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002640
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002641 for (i=0 ; i<r ; i++)
2642 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002643
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002644 /* create cwrobject structure */
2645 co = (cwrobject *)type->tp_alloc(type, 0);
2646 if (co == NULL)
2647 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002648
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002649 co->pool = pool;
2650 co->indices = indices;
2651 co->result = NULL;
2652 co->r = r;
2653 co->stopped = !n && r;
2654
2655 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002656
2657error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002658 if (indices != NULL)
2659 PyMem_Free(indices);
2660 Py_XDECREF(pool);
2661 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002662}
2663
2664static void
2665cwr_dealloc(cwrobject *co)
2666{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002667 PyObject_GC_UnTrack(co);
2668 Py_XDECREF(co->pool);
2669 Py_XDECREF(co->result);
2670 if (co->indices != NULL)
2671 PyMem_Free(co->indices);
2672 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002673}
2674
2675static int
2676cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2677{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002678 Py_VISIT(co->pool);
2679 Py_VISIT(co->result);
2680 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002681}
2682
2683static PyObject *
2684cwr_next(cwrobject *co)
2685{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002686 PyObject *elem;
2687 PyObject *oldelem;
2688 PyObject *pool = co->pool;
2689 Py_ssize_t *indices = co->indices;
2690 PyObject *result = co->result;
2691 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2692 Py_ssize_t r = co->r;
2693 Py_ssize_t i, j, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002694
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002695 if (co->stopped)
2696 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002697
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 if (result == NULL) {
2699 /* On the first pass, initialize result tuple using the indices */
2700 result = PyTuple_New(r);
2701 if (result == NULL)
2702 goto empty;
2703 co->result = result;
2704 for (i=0; i<r ; i++) {
2705 index = indices[i];
2706 elem = PyTuple_GET_ITEM(pool, index);
2707 Py_INCREF(elem);
2708 PyTuple_SET_ITEM(result, i, elem);
2709 }
2710 } else {
2711 /* Copy the previous result tuple or re-use it if available */
2712 if (Py_REFCNT(result) > 1) {
2713 PyObject *old_result = result;
2714 result = PyTuple_New(r);
2715 if (result == NULL)
2716 goto empty;
2717 co->result = result;
2718 for (i=0; i<r ; i++) {
2719 elem = PyTuple_GET_ITEM(old_result, i);
2720 Py_INCREF(elem);
2721 PyTuple_SET_ITEM(result, i, elem);
2722 }
2723 Py_DECREF(old_result);
2724 }
2725 /* Now, we've got the only copy so we can update it in-place CPython's
2726 empty tuple is a singleton and cached in PyTuple's freelist. */
2727 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002728
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002729 /* Scan indices right-to-left until finding one that is not
2730 * at its maximum (n-1). */
2731 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2732 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002733
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002734 /* If i is negative, then the indices are all at
2735 their maximum value and we're done. */
2736 if (i < 0)
2737 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002738
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002739 /* Increment the current index which we know is not at its
2740 maximum. Then set all to the right to the same value. */
2741 indices[i]++;
2742 for (j=i+1 ; j<r ; j++)
2743 indices[j] = indices[j-1];
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002745 /* Update the result tuple for the new indices
2746 starting with i, the leftmost index that changed */
2747 for ( ; i<r ; i++) {
2748 index = indices[i];
2749 elem = PyTuple_GET_ITEM(pool, index);
2750 Py_INCREF(elem);
2751 oldelem = PyTuple_GET_ITEM(result, i);
2752 PyTuple_SET_ITEM(result, i, elem);
2753 Py_DECREF(oldelem);
2754 }
2755 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002757 Py_INCREF(result);
2758 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002759
2760empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002761 co->stopped = 1;
2762 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002763}
2764
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002765static PyObject *
2766cwr_reduce(cwrobject *lz)
2767{
2768 if (lz->result == NULL) {
2769 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2770 } else if (lz->stopped) {
2771 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2772 } else {
2773 PyObject *indices;
2774 Py_ssize_t i;
2775
2776 /* we must pickle the indices and use them for setstate */
2777 indices = PyTuple_New(lz->r);
2778 if (!indices)
2779 return NULL;
2780 for (i=0; i<lz->r; i++)
2781 {
2782 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2783 if (!index) {
2784 Py_DECREF(indices);
2785 return NULL;
2786 }
2787 PyTuple_SET_ITEM(indices, i, index);
2788 }
2789
2790 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2791 }
2792}
2793
2794static PyObject *
2795cwr_setstate(cwrobject *lz, PyObject *state)
2796{
2797 PyObject *result;
2798 Py_ssize_t n, i;
2799
2800 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2801 {
2802 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2803 return NULL;
2804 }
2805
2806 n = PyTuple_GET_SIZE(lz->pool);
2807 for (i=0; i<lz->r; i++)
2808 {
2809 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2810 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2811 if (index < 0 && PyErr_Occurred())
2812 return NULL; /* not an integer */
2813 /* clamp the index */
2814 if (index < 0)
2815 index = 0;
2816 else if (index > n-1)
2817 index = n-1;
2818 lz->indices[i] = index;
2819 }
2820 result = PyTuple_New(lz->r);
2821 if (result == NULL)
2822 return NULL;
2823 for (i=0; i<lz->r; i++) {
2824 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2825 Py_INCREF(element);
2826 PyTuple_SET_ITEM(result, i, element);
2827 }
2828 Py_CLEAR(lz->result);
2829 lz->result = result;
2830 Py_RETURN_NONE;
2831}
2832
2833static PyMethodDef cwr_methods[] = {
2834 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
2835 reduce_doc},
2836 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
2837 setstate_doc},
2838 {NULL, NULL} /* sentinel */
2839};
2840
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002841PyDoc_STRVAR(cwr_doc,
Raymond Hettinger36c3c022009-11-19 01:20:07 +00002842"combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002843\n\
2844Return successive r-length combinations of elements in the iterable\n\
2845allowing individual elements to have successive repeats.\n\
2846combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2847
2848static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002849 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002850 "itertools.combinations_with_replacement", /* tp_name */
2851 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002852 0, /* tp_itemsize */
2853 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002854 (destructor)cwr_dealloc, /* tp_dealloc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002855 0, /* tp_print */
2856 0, /* tp_getattr */
2857 0, /* tp_setattr */
2858 0, /* tp_reserved */
2859 0, /* tp_repr */
2860 0, /* tp_as_number */
2861 0, /* tp_as_sequence */
2862 0, /* tp_as_mapping */
2863 0, /* tp_hash */
2864 0, /* tp_call */
2865 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002866 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002867 0, /* tp_setattro */
2868 0, /* tp_as_buffer */
2869 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002870 Py_TPFLAGS_BASETYPE, /* tp_flags */
2871 cwr_doc, /* tp_doc */
2872 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002873 0, /* tp_clear */
2874 0, /* tp_richcompare */
2875 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002876 PyObject_SelfIter, /* tp_iter */
2877 (iternextfunc)cwr_next, /* tp_iternext */
2878 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002879 0, /* tp_members */
2880 0, /* tp_getset */
2881 0, /* tp_base */
2882 0, /* tp_dict */
2883 0, /* tp_descr_get */
2884 0, /* tp_descr_set */
2885 0, /* tp_dictoffset */
2886 0, /* tp_init */
2887 0, /* tp_alloc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002888 cwr_new, /* tp_new */
2889 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002890};
2891
2892
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002893/* permutations object ************************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002894
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002895def permutations(iterable, r=None):
2896 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2897 pool = tuple(iterable)
2898 n = len(pool)
2899 r = n if r is None else r
2900 indices = range(n)
2901 cycles = range(n-r+1, n+1)[::-1]
2902 yield tuple(pool[i] for i in indices[:r])
2903 while n:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002904 for i in reversed(range(r)):
2905 cycles[i] -= 1
2906 if cycles[i] == 0:
2907 indices[i:] = indices[i+1:] + indices[i:i+1]
2908 cycles[i] = n - i
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002909 else:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002910 j = cycles[i]
2911 indices[i], indices[-j] = indices[-j], indices[i]
2912 yield tuple(pool[i] for i in indices[:r])
2913 break
2914 else:
2915 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002916*/
2917
2918typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002919 PyObject_HEAD
2920 PyObject *pool; /* input converted to a tuple */
2921 Py_ssize_t *indices; /* one index per element in the pool */
2922 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2923 PyObject *result; /* most recently returned result tuple */
2924 Py_ssize_t r; /* size of result tuple */
2925 int stopped; /* set to 1 when the permutations iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002926} permutationsobject;
2927
2928static PyTypeObject permutations_type;
2929
2930static PyObject *
2931permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2932{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 permutationsobject *po;
2934 Py_ssize_t n;
2935 Py_ssize_t r;
2936 PyObject *robj = Py_None;
2937 PyObject *pool = NULL;
2938 PyObject *iterable = NULL;
2939 Py_ssize_t *indices = NULL;
2940 Py_ssize_t *cycles = NULL;
2941 Py_ssize_t i;
2942 static char *kwargs[] = {"iterable", "r", NULL};
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002943
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002944 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2945 &iterable, &robj))
2946 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002947
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002948 pool = PySequence_Tuple(iterable);
2949 if (pool == NULL)
2950 goto error;
2951 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002952
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002953 r = n;
2954 if (robj != Py_None) {
2955 if (!PyLong_Check(robj)) {
2956 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2957 goto error;
2958 }
2959 r = PyLong_AsSsize_t(robj);
2960 if (r == -1 && PyErr_Occurred())
2961 goto error;
2962 }
2963 if (r < 0) {
2964 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2965 goto error;
2966 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002967
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002968 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2969 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2970 if (indices == NULL || cycles == NULL) {
2971 PyErr_NoMemory();
2972 goto error;
2973 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002974
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002975 for (i=0 ; i<n ; i++)
2976 indices[i] = i;
2977 for (i=0 ; i<r ; i++)
2978 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002979
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002980 /* create permutationsobject structure */
2981 po = (permutationsobject *)type->tp_alloc(type, 0);
2982 if (po == NULL)
2983 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002984
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002985 po->pool = pool;
2986 po->indices = indices;
2987 po->cycles = cycles;
2988 po->result = NULL;
2989 po->r = r;
2990 po->stopped = r > n ? 1 : 0;
2991
2992 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00002993
2994error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002995 if (indices != NULL)
2996 PyMem_Free(indices);
2997 if (cycles != NULL)
2998 PyMem_Free(cycles);
2999 Py_XDECREF(pool);
3000 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003001}
3002
3003static void
3004permutations_dealloc(permutationsobject *po)
3005{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003006 PyObject_GC_UnTrack(po);
3007 Py_XDECREF(po->pool);
3008 Py_XDECREF(po->result);
3009 PyMem_Free(po->indices);
3010 PyMem_Free(po->cycles);
3011 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003012}
3013
3014static int
3015permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3016{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003017 Py_VISIT(po->pool);
3018 Py_VISIT(po->result);
3019 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003020}
3021
3022static PyObject *
3023permutations_next(permutationsobject *po)
3024{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 PyObject *elem;
3026 PyObject *oldelem;
3027 PyObject *pool = po->pool;
3028 Py_ssize_t *indices = po->indices;
3029 Py_ssize_t *cycles = po->cycles;
3030 PyObject *result = po->result;
3031 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3032 Py_ssize_t r = po->r;
3033 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003034
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003035 if (po->stopped)
3036 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003037
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003038 if (result == NULL) {
3039 /* On the first pass, initialize result tuple using the indices */
3040 result = PyTuple_New(r);
3041 if (result == NULL)
3042 goto empty;
3043 po->result = result;
3044 for (i=0; i<r ; i++) {
3045 index = indices[i];
3046 elem = PyTuple_GET_ITEM(pool, index);
3047 Py_INCREF(elem);
3048 PyTuple_SET_ITEM(result, i, elem);
3049 }
3050 } else {
3051 if (n == 0)
3052 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003054 /* Copy the previous result tuple or re-use it if available */
3055 if (Py_REFCNT(result) > 1) {
3056 PyObject *old_result = result;
3057 result = PyTuple_New(r);
3058 if (result == NULL)
3059 goto empty;
3060 po->result = result;
3061 for (i=0; i<r ; i++) {
3062 elem = PyTuple_GET_ITEM(old_result, i);
3063 Py_INCREF(elem);
3064 PyTuple_SET_ITEM(result, i, elem);
3065 }
3066 Py_DECREF(old_result);
3067 }
3068 /* Now, we've got the only copy so we can update it in-place */
3069 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003070
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003071 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3072 for (i=r-1 ; i>=0 ; i--) {
3073 cycles[i] -= 1;
3074 if (cycles[i] == 0) {
3075 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3076 index = indices[i];
3077 for (j=i ; j<n-1 ; j++)
3078 indices[j] = indices[j+1];
3079 indices[n-1] = index;
3080 cycles[i] = n - i;
3081 } else {
3082 j = cycles[i];
3083 index = indices[i];
3084 indices[i] = indices[n-j];
3085 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003086
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003087 for (k=i; k<r ; k++) {
3088 /* start with i, the leftmost element that changed */
3089 /* yield tuple(pool[k] for k in indices[:r]) */
3090 index = indices[k];
3091 elem = PyTuple_GET_ITEM(pool, index);
3092 Py_INCREF(elem);
3093 oldelem = PyTuple_GET_ITEM(result, k);
3094 PyTuple_SET_ITEM(result, k, elem);
3095 Py_DECREF(oldelem);
3096 }
3097 break;
3098 }
3099 }
3100 /* If i is negative, then the cycles have all
3101 rolled-over and we're done. */
3102 if (i < 0)
3103 goto empty;
3104 }
3105 Py_INCREF(result);
3106 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003107
3108empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003109 po->stopped = 1;
3110 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003111}
3112
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003113static PyObject *
3114permutations_reduce(permutationsobject *po)
3115{
3116 if (po->result == NULL) {
3117 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3118 } else if (po->stopped) {
3119 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3120 } else {
3121 PyObject *indices=NULL, *cycles=NULL;
3122 Py_ssize_t n, i;
3123
3124 /* we must pickle the indices and cycles and use them for setstate */
3125 n = PyTuple_GET_SIZE(po->pool);
3126 indices = PyTuple_New(n);
3127 if (indices == NULL)
3128 goto err;
3129 for (i=0; i<n; i++){
3130 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3131 if (!index)
3132 goto err;
3133 PyTuple_SET_ITEM(indices, i, index);
3134 }
3135
3136 cycles = PyTuple_New(po->r);
3137 if (cycles == NULL)
3138 goto err;
3139 for (i=0; i<po->r; i++)
3140 {
3141 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3142 if (!index)
3143 goto err;
3144 PyTuple_SET_ITEM(cycles, i, index);
3145 }
3146 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3147 po->pool, po->r,
3148 indices, cycles);
3149 err:
3150 Py_XDECREF(indices);
3151 Py_XDECREF(cycles);
3152 return NULL;
3153 }
3154}
3155
3156static PyObject *
3157permutations_setstate(permutationsobject *po, PyObject *state)
3158{
3159 PyObject *indices, *cycles, *result;
3160 Py_ssize_t n, i;
3161
3162 if (!PyArg_ParseTuple(state, "O!O!",
3163 &PyTuple_Type, &indices,
3164 &PyTuple_Type, &cycles))
3165 return NULL;
3166
3167 n = PyTuple_GET_SIZE(po->pool);
3168 if (PyTuple_GET_SIZE(indices) != n ||
3169 PyTuple_GET_SIZE(cycles) != po->r)
3170 {
3171 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3172 return NULL;
3173 }
3174
3175 for (i=0; i<n; i++)
3176 {
3177 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3178 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3179 if (index < 0 && PyErr_Occurred())
3180 return NULL; /* not an integer */
3181 /* clamp the index */
3182 if (index < 0)
3183 index = 0;
3184 else if (index > n-1)
3185 index = n-1;
3186 po->indices[i] = index;
3187 }
3188
3189 for (i=0; i<po->r; i++)
3190 {
3191 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3192 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3193 if (index < 0 && PyErr_Occurred())
3194 return NULL; /* not an integer */
3195 if (index < 1)
3196 index = 1;
3197 else if (index > n-i)
3198 index = n-i;
3199 po->cycles[i] = index;
3200 }
3201 result = PyTuple_New(po->r);
3202 if (result == NULL)
3203 return NULL;
3204 for (i=0; i<po->r; i++) {
3205 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3206 Py_INCREF(element);
3207 PyTuple_SET_ITEM(result, i, element);
3208 }
3209 Py_CLEAR(po->result);
3210 po->result = result;
3211 Py_RETURN_NONE;
3212}
3213
3214static PyMethodDef permuations_methods[] = {
3215 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3216 reduce_doc},
3217 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3218 setstate_doc},
3219 {NULL, NULL} /* sentinel */
3220};
3221
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003222PyDoc_STRVAR(permutations_doc,
Georg Brandl0c77a822008-06-10 16:37:50 +00003223"permutations(iterable[, r]) --> permutations object\n\
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003224\n\
3225Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl0c77a822008-06-10 16:37:50 +00003226permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003227
3228static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003229 PyVarObject_HEAD_INIT(NULL, 0)
3230 "itertools.permutations", /* tp_name */
3231 sizeof(permutationsobject), /* tp_basicsize */
3232 0, /* tp_itemsize */
3233 /* methods */
3234 (destructor)permutations_dealloc, /* tp_dealloc */
3235 0, /* tp_print */
3236 0, /* tp_getattr */
3237 0, /* tp_setattr */
3238 0, /* tp_reserved */
3239 0, /* tp_repr */
3240 0, /* tp_as_number */
3241 0, /* tp_as_sequence */
3242 0, /* tp_as_mapping */
3243 0, /* tp_hash */
3244 0, /* tp_call */
3245 0, /* tp_str */
3246 PyObject_GenericGetAttr, /* tp_getattro */
3247 0, /* tp_setattro */
3248 0, /* tp_as_buffer */
3249 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3250 Py_TPFLAGS_BASETYPE, /* tp_flags */
3251 permutations_doc, /* tp_doc */
3252 (traverseproc)permutations_traverse, /* tp_traverse */
3253 0, /* tp_clear */
3254 0, /* tp_richcompare */
3255 0, /* tp_weaklistoffset */
3256 PyObject_SelfIter, /* tp_iter */
3257 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003258 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003259 0, /* tp_members */
3260 0, /* tp_getset */
3261 0, /* tp_base */
3262 0, /* tp_dict */
3263 0, /* tp_descr_get */
3264 0, /* tp_descr_set */
3265 0, /* tp_dictoffset */
3266 0, /* tp_init */
3267 0, /* tp_alloc */
3268 permutations_new, /* tp_new */
3269 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003270};
3271
Raymond Hettinger482ba772010-12-01 22:48:00 +00003272/* accumulate object ************************************************************/
3273
3274typedef struct {
3275 PyObject_HEAD
3276 PyObject *total;
3277 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003278 PyObject *binop;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003279} accumulateobject;
3280
3281static PyTypeObject accumulate_type;
3282
3283static PyObject *
3284accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3285{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003286 static char *kwargs[] = {"iterable", "func", NULL};
Raymond Hettinger482ba772010-12-01 22:48:00 +00003287 PyObject *iterable;
3288 PyObject *it;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003289 PyObject *binop = Py_None;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003290 accumulateobject *lz;
3291
Raymond Hettinger5d446132011-03-27 18:52:10 -07003292 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate",
3293 kwargs, &iterable, &binop))
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003294 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003295
3296 /* Get iterator. */
3297 it = PyObject_GetIter(iterable);
3298 if (it == NULL)
3299 return NULL;
3300
Raymond Hettinger482ba772010-12-01 22:48:00 +00003301 /* create accumulateobject structure */
3302 lz = (accumulateobject *)type->tp_alloc(type, 0);
3303 if (lz == NULL) {
3304 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003305 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003306 }
3307
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003308 if (binop != Py_None) {
3309 Py_XINCREF(binop);
3310 lz->binop = binop;
3311 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003312 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003313 lz->it = it;
3314 return (PyObject *)lz;
3315}
3316
3317static void
3318accumulate_dealloc(accumulateobject *lz)
3319{
3320 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003321 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003322 Py_XDECREF(lz->total);
3323 Py_XDECREF(lz->it);
3324 Py_TYPE(lz)->tp_free(lz);
3325}
3326
3327static int
3328accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3329{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003330 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003331 Py_VISIT(lz->it);
3332 Py_VISIT(lz->total);
3333 return 0;
3334}
3335
3336static PyObject *
3337accumulate_next(accumulateobject *lz)
3338{
3339 PyObject *val, *oldtotal, *newtotal;
3340
3341 val = PyIter_Next(lz->it);
3342 if (val == NULL)
3343 return NULL;
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003344
3345 if (lz->total == NULL) {
3346 Py_INCREF(val);
3347 lz->total = val;
3348 return lz->total;
3349 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003350
3351 if (lz->binop == NULL)
3352 newtotal = PyNumber_Add(lz->total, val);
3353 else
3354 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003355 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003356 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003357 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003358
3359 oldtotal = lz->total;
3360 lz->total = newtotal;
3361 Py_DECREF(oldtotal);
3362
3363 Py_INCREF(newtotal);
3364 return newtotal;
3365}
3366
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003367static PyObject *
3368accumulate_reduce(accumulateobject *lz)
3369{
3370 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3371 lz->it, lz->binop?lz->binop:Py_None,
3372 lz->total?lz->total:Py_None);
3373 }
3374
3375static PyObject *
3376accumulate_setstate(accumulateobject *lz, PyObject *state)
3377{
3378 Py_CLEAR(lz->total);
3379 lz->total = state;
3380 Py_INCREF(lz->total);
3381 Py_RETURN_NONE;
3382}
3383
3384static PyMethodDef accumulate_methods[] = {
3385 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3386 reduce_doc},
3387 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3388 setstate_doc},
3389 {NULL, NULL} /* sentinel */
3390};
3391
Raymond Hettinger482ba772010-12-01 22:48:00 +00003392PyDoc_STRVAR(accumulate_doc,
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003393"accumulate(iterable[, func]) --> accumulate object\n\
Raymond Hettinger482ba772010-12-01 22:48:00 +00003394\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07003395Return series of accumulated sums (or other binary function results).");
Raymond Hettinger482ba772010-12-01 22:48:00 +00003396
3397static PyTypeObject accumulate_type = {
3398 PyVarObject_HEAD_INIT(NULL, 0)
3399 "itertools.accumulate", /* tp_name */
3400 sizeof(accumulateobject), /* tp_basicsize */
3401 0, /* tp_itemsize */
3402 /* methods */
3403 (destructor)accumulate_dealloc, /* tp_dealloc */
3404 0, /* tp_print */
3405 0, /* tp_getattr */
3406 0, /* tp_setattr */
3407 0, /* tp_reserved */
3408 0, /* tp_repr */
3409 0, /* tp_as_number */
3410 0, /* tp_as_sequence */
3411 0, /* tp_as_mapping */
3412 0, /* tp_hash */
3413 0, /* tp_call */
3414 0, /* tp_str */
3415 PyObject_GenericGetAttr, /* tp_getattro */
3416 0, /* tp_setattro */
3417 0, /* tp_as_buffer */
3418 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3419 Py_TPFLAGS_BASETYPE, /* tp_flags */
3420 accumulate_doc, /* tp_doc */
3421 (traverseproc)accumulate_traverse, /* tp_traverse */
3422 0, /* tp_clear */
3423 0, /* tp_richcompare */
3424 0, /* tp_weaklistoffset */
3425 PyObject_SelfIter, /* tp_iter */
3426 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003427 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003428 0, /* tp_members */
3429 0, /* tp_getset */
3430 0, /* tp_base */
3431 0, /* tp_dict */
3432 0, /* tp_descr_get */
3433 0, /* tp_descr_set */
3434 0, /* tp_dictoffset */
3435 0, /* tp_init */
3436 0, /* tp_alloc */
3437 accumulate_new, /* tp_new */
3438 PyObject_GC_Del, /* tp_free */
3439};
3440
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003441
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003442/* compress object ************************************************************/
3443
3444/* Equivalent to:
3445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003446 def compress(data, selectors):
3447 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3448 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003449*/
3450
3451typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003452 PyObject_HEAD
3453 PyObject *data;
3454 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003455} compressobject;
3456
3457static PyTypeObject compress_type;
3458
3459static PyObject *
3460compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3461{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003462 PyObject *seq1, *seq2;
3463 PyObject *data=NULL, *selectors=NULL;
3464 compressobject *lz;
3465 static char *kwargs[] = {"data", "selectors", NULL};
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003467 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
3468 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003469
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003470 data = PyObject_GetIter(seq1);
3471 if (data == NULL)
3472 goto fail;
3473 selectors = PyObject_GetIter(seq2);
3474 if (selectors == NULL)
3475 goto fail;
3476
3477 /* create compressobject structure */
3478 lz = (compressobject *)type->tp_alloc(type, 0);
3479 if (lz == NULL)
3480 goto fail;
3481 lz->data = data;
3482 lz->selectors = selectors;
3483 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003484
3485fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003486 Py_XDECREF(data);
3487 Py_XDECREF(selectors);
3488 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003489}
3490
3491static void
3492compress_dealloc(compressobject *lz)
3493{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003494 PyObject_GC_UnTrack(lz);
3495 Py_XDECREF(lz->data);
3496 Py_XDECREF(lz->selectors);
3497 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003498}
3499
3500static int
3501compress_traverse(compressobject *lz, visitproc visit, void *arg)
3502{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003503 Py_VISIT(lz->data);
3504 Py_VISIT(lz->selectors);
3505 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003506}
3507
3508static PyObject *
3509compress_next(compressobject *lz)
3510{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003511 PyObject *data = lz->data, *selectors = lz->selectors;
3512 PyObject *datum, *selector;
3513 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3514 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3515 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003516
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003517 while (1) {
3518 /* Steps: get datum, get selector, evaluate selector.
3519 Order is important (to match the pure python version
3520 in terms of which input gets a chance to raise an
3521 exception first).
3522 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003523
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003524 datum = datanext(data);
3525 if (datum == NULL)
3526 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003527
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003528 selector = selectornext(selectors);
3529 if (selector == NULL) {
3530 Py_DECREF(datum);
3531 return NULL;
3532 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003533
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003534 ok = PyObject_IsTrue(selector);
3535 Py_DECREF(selector);
3536 if (ok == 1)
3537 return datum;
3538 Py_DECREF(datum);
3539 if (ok == -1)
3540 return NULL;
3541 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003542}
3543
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003544static PyObject *
3545compress_reduce(compressobject *lz)
3546{
3547 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3548 lz->data, lz->selectors);
3549 }
3550
3551static PyMethodDef compress_methods[] = {
3552 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3553 reduce_doc},
3554 {NULL, NULL} /* sentinel */
3555};
3556
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003557PyDoc_STRVAR(compress_doc,
Raymond Hettinger15a49502009-02-19 02:17:09 +00003558"compress(data, selectors) --> iterator over selected data\n\
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003559\n\
3560Return data elements corresponding to true selector elements.\n\
3561Forms a shorter iterator from selected data elements using the\n\
3562selectors to choose the data elements.");
3563
3564static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003565 PyVarObject_HEAD_INIT(NULL, 0)
3566 "itertools.compress", /* tp_name */
3567 sizeof(compressobject), /* tp_basicsize */
3568 0, /* tp_itemsize */
3569 /* methods */
3570 (destructor)compress_dealloc, /* tp_dealloc */
3571 0, /* tp_print */
3572 0, /* tp_getattr */
3573 0, /* tp_setattr */
3574 0, /* tp_reserved */
3575 0, /* tp_repr */
3576 0, /* tp_as_number */
3577 0, /* tp_as_sequence */
3578 0, /* tp_as_mapping */
3579 0, /* tp_hash */
3580 0, /* tp_call */
3581 0, /* tp_str */
3582 PyObject_GenericGetAttr, /* tp_getattro */
3583 0, /* tp_setattro */
3584 0, /* tp_as_buffer */
3585 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3586 Py_TPFLAGS_BASETYPE, /* tp_flags */
3587 compress_doc, /* tp_doc */
3588 (traverseproc)compress_traverse, /* tp_traverse */
3589 0, /* tp_clear */
3590 0, /* tp_richcompare */
3591 0, /* tp_weaklistoffset */
3592 PyObject_SelfIter, /* tp_iter */
3593 (iternextfunc)compress_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003594 compress_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003595 0, /* tp_members */
3596 0, /* tp_getset */
3597 0, /* tp_base */
3598 0, /* tp_dict */
3599 0, /* tp_descr_get */
3600 0, /* tp_descr_set */
3601 0, /* tp_dictoffset */
3602 0, /* tp_init */
3603 0, /* tp_alloc */
3604 compress_new, /* tp_new */
3605 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003606};
3607
3608
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003609/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003610
3611typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003612 PyObject_HEAD
3613 PyObject *func;
3614 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003615} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003616
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003617static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003618
3619static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003620filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003621{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003622 PyObject *func, *seq;
3623 PyObject *it;
3624 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003625
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003626 if (type == &filterfalse_type &&
3627 !_PyArg_NoKeywords("filterfalse()", kwds))
3628 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003629
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003630 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
3631 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003632
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003633 /* Get iterator. */
3634 it = PyObject_GetIter(seq);
3635 if (it == NULL)
3636 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003637
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003638 /* create filterfalseobject structure */
3639 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3640 if (lz == NULL) {
3641 Py_DECREF(it);
3642 return NULL;
3643 }
3644 Py_INCREF(func);
3645 lz->func = func;
3646 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003647
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003648 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003649}
3650
3651static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003652filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003653{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003654 PyObject_GC_UnTrack(lz);
3655 Py_XDECREF(lz->func);
3656 Py_XDECREF(lz->it);
3657 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003658}
3659
3660static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003661filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003662{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003663 Py_VISIT(lz->it);
3664 Py_VISIT(lz->func);
3665 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003666}
3667
3668static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003669filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003670{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003671 PyObject *item;
3672 PyObject *it = lz->it;
3673 long ok;
3674 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003675
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003676 iternext = *Py_TYPE(it)->tp_iternext;
3677 for (;;) {
3678 item = iternext(it);
3679 if (item == NULL)
3680 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003681
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003682 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3683 ok = PyObject_IsTrue(item);
3684 } else {
3685 PyObject *good;
3686 good = PyObject_CallFunctionObjArgs(lz->func,
3687 item, NULL);
3688 if (good == NULL) {
3689 Py_DECREF(item);
3690 return NULL;
3691 }
3692 ok = PyObject_IsTrue(good);
3693 Py_DECREF(good);
3694 }
3695 if (!ok)
3696 return item;
3697 Py_DECREF(item);
3698 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003699}
3700
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003701static PyObject *
3702filterfalse_reduce(filterfalseobject *lz)
3703{
3704 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3705 lz->func, lz->it);
3706 }
3707
3708static PyMethodDef filterfalse_methods[] = {
3709 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3710 reduce_doc},
3711 {NULL, NULL} /* sentinel */
3712};
3713
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003714PyDoc_STRVAR(filterfalse_doc,
3715"filterfalse(function or None, sequence) --> filterfalse object\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003716\n\
3717Return those items of sequence for which function(item) is false.\n\
3718If function is None, return the items that are false.");
3719
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003720static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003721 PyVarObject_HEAD_INIT(NULL, 0)
3722 "itertools.filterfalse", /* tp_name */
3723 sizeof(filterfalseobject), /* tp_basicsize */
3724 0, /* tp_itemsize */
3725 /* methods */
3726 (destructor)filterfalse_dealloc, /* tp_dealloc */
3727 0, /* tp_print */
3728 0, /* tp_getattr */
3729 0, /* tp_setattr */
3730 0, /* tp_reserved */
3731 0, /* tp_repr */
3732 0, /* tp_as_number */
3733 0, /* tp_as_sequence */
3734 0, /* tp_as_mapping */
3735 0, /* tp_hash */
3736 0, /* tp_call */
3737 0, /* tp_str */
3738 PyObject_GenericGetAttr, /* tp_getattro */
3739 0, /* tp_setattro */
3740 0, /* tp_as_buffer */
3741 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3742 Py_TPFLAGS_BASETYPE, /* tp_flags */
3743 filterfalse_doc, /* tp_doc */
3744 (traverseproc)filterfalse_traverse, /* tp_traverse */
3745 0, /* tp_clear */
3746 0, /* tp_richcompare */
3747 0, /* tp_weaklistoffset */
3748 PyObject_SelfIter, /* tp_iter */
3749 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003750 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003751 0, /* tp_members */
3752 0, /* tp_getset */
3753 0, /* tp_base */
3754 0, /* tp_dict */
3755 0, /* tp_descr_get */
3756 0, /* tp_descr_set */
3757 0, /* tp_dictoffset */
3758 0, /* tp_init */
3759 0, /* tp_alloc */
3760 filterfalse_new, /* tp_new */
3761 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003762};
3763
3764
3765/* count object ************************************************************/
3766
3767typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003768 PyObject_HEAD
3769 Py_ssize_t cnt;
3770 PyObject *long_cnt;
3771 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003772} countobject;
3773
Raymond Hettinger30729212009-02-12 06:28:27 +00003774/* Counting logic and invariants:
3775
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003776fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003777
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003778 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3779 Advances with: cnt += 1
3780 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003781
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003782slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003783
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003784 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3785 All counting is done with python objects (no overflows or underflows).
3786 Advances with: long_cnt += long_step
3787 Step may be zero -- effectively a slow version of repeat(cnt).
3788 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003789*/
3790
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003791static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003792
3793static PyObject *
3794count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3795{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003796 countobject *lz;
3797 int slow_mode = 0;
3798 Py_ssize_t cnt = 0;
3799 PyObject *long_cnt = NULL;
3800 PyObject *long_step = NULL;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003801 long step;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003802 static char *kwlist[] = {"start", "step", 0};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003803
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003804 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3805 kwlist, &long_cnt, &long_step))
3806 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003807
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003808 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3809 (long_step != NULL && !PyNumber_Check(long_step))) {
3810 PyErr_SetString(PyExc_TypeError, "a number is required");
3811 return NULL;
3812 }
Raymond Hettinger30729212009-02-12 06:28:27 +00003813
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003814 if (long_cnt != NULL) {
3815 cnt = PyLong_AsSsize_t(long_cnt);
3816 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
3817 PyErr_Clear();
3818 slow_mode = 1;
3819 }
3820 Py_INCREF(long_cnt);
3821 } else {
3822 cnt = 0;
3823 long_cnt = PyLong_FromLong(0);
3824 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003825
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003826 /* If not specified, step defaults to 1 */
3827 if (long_step == NULL) {
3828 long_step = PyLong_FromLong(1);
3829 if (long_step == NULL) {
3830 Py_DECREF(long_cnt);
3831 return NULL;
3832 }
3833 } else
3834 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003835
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003836 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003837
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003838 /* Fast mode only works when the step is 1 */
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00003839 step = PyLong_AsLong(long_step);
3840 if (step != 1) {
3841 slow_mode = 1;
3842 if (step == -1 && PyErr_Occurred())
3843 PyErr_Clear();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003844 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003845
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003846 if (slow_mode)
3847 cnt = PY_SSIZE_T_MAX;
3848 else
3849 Py_CLEAR(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003850
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003851 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3852 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3853 assert(slow_mode ||
3854 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003855
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003856 /* create countobject structure */
3857 lz = (countobject *)type->tp_alloc(type, 0);
3858 if (lz == NULL) {
3859 Py_XDECREF(long_cnt);
3860 return NULL;
3861 }
3862 lz->cnt = cnt;
3863 lz->long_cnt = long_cnt;
3864 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003865
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003866 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003867}
3868
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003869static void
3870count_dealloc(countobject *lz)
3871{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003872 PyObject_GC_UnTrack(lz);
3873 Py_XDECREF(lz->long_cnt);
3874 Py_XDECREF(lz->long_step);
3875 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003876}
3877
Benjamin Peterson25e26a02009-02-16 21:28:29 +00003878static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00003879count_traverse(countobject *lz, visitproc visit, void *arg)
3880{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003881 Py_VISIT(lz->long_cnt);
3882 Py_VISIT(lz->long_step);
3883 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003884}
3885
3886static PyObject *
3887count_nextlong(countobject *lz)
3888{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003889 PyObject *long_cnt;
3890 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003891
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003892 long_cnt = lz->long_cnt;
3893 if (long_cnt == NULL) {
3894 /* Switch to slow_mode */
3895 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3896 if (long_cnt == NULL)
3897 return NULL;
3898 }
3899 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00003900
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003901 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3902 if (stepped_up == NULL)
3903 return NULL;
3904 lz->long_cnt = stepped_up;
3905 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003906}
3907
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003908static PyObject *
3909count_next(countobject *lz)
3910{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003911 if (lz->cnt == PY_SSIZE_T_MAX)
3912 return count_nextlong(lz);
3913 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003914}
3915
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003916static PyObject *
3917count_repr(countobject *lz)
3918{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003919 if (lz->cnt != PY_SSIZE_T_MAX)
3920 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00003921
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003922 if (PyLong_Check(lz->long_step)) {
3923 long step = PyLong_AsLong(lz->long_step);
3924 if (step == -1 && PyErr_Occurred()) {
3925 PyErr_Clear();
3926 }
3927 if (step == 1) {
3928 /* Don't display step when it is an integer equal to 1 */
3929 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3930 }
3931 }
3932 return PyUnicode_FromFormat("count(%R, %R)",
3933 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003934}
3935
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003936static PyObject *
3937count_reduce(countobject *lz)
3938{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003939 if (lz->cnt == PY_SSIZE_T_MAX)
3940 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3941 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003942}
3943
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003944static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003945 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003946 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00003948};
3949
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003950PyDoc_STRVAR(count_doc,
Ezio Melottibfd73fa2010-06-11 02:26:42 +00003951 "count(start=0, step=1) --> count object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003952\n\
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003953Return a count object whose .__next__() method returns consecutive values.\n\
3954Equivalent to:\n\n\
Raymond Hettinger30729212009-02-12 06:28:27 +00003955 def count(firstval=0, step=1):\n\
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003956 x = firstval\n\
3957 while 1:\n\
3958 yield x\n\
3959 x += step\n");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003960
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003961static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003962 PyVarObject_HEAD_INIT(NULL, 0)
3963 "itertools.count", /* tp_name */
3964 sizeof(countobject), /* tp_basicsize */
3965 0, /* tp_itemsize */
3966 /* methods */
3967 (destructor)count_dealloc, /* tp_dealloc */
3968 0, /* tp_print */
3969 0, /* tp_getattr */
3970 0, /* tp_setattr */
3971 0, /* tp_reserved */
3972 (reprfunc)count_repr, /* tp_repr */
3973 0, /* tp_as_number */
3974 0, /* tp_as_sequence */
3975 0, /* tp_as_mapping */
3976 0, /* tp_hash */
3977 0, /* tp_call */
3978 0, /* tp_str */
3979 PyObject_GenericGetAttr, /* tp_getattro */
3980 0, /* tp_setattro */
3981 0, /* tp_as_buffer */
3982 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3983 Py_TPFLAGS_BASETYPE, /* tp_flags */
3984 count_doc, /* tp_doc */
3985 (traverseproc)count_traverse, /* tp_traverse */
3986 0, /* tp_clear */
3987 0, /* tp_richcompare */
3988 0, /* tp_weaklistoffset */
3989 PyObject_SelfIter, /* tp_iter */
3990 (iternextfunc)count_next, /* tp_iternext */
3991 count_methods, /* tp_methods */
3992 0, /* tp_members */
3993 0, /* tp_getset */
3994 0, /* tp_base */
3995 0, /* tp_dict */
3996 0, /* tp_descr_get */
3997 0, /* tp_descr_set */
3998 0, /* tp_dictoffset */
3999 0, /* tp_init */
4000 0, /* tp_alloc */
4001 count_new, /* tp_new */
4002 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004003};
4004
4005
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004006/* repeat object ************************************************************/
4007
4008typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004009 PyObject_HEAD
4010 PyObject *element;
4011 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004012} repeatobject;
4013
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004014static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004015
4016static PyObject *
4017repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4018{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004019 repeatobject *ro;
4020 PyObject *element;
4021 Py_ssize_t cnt = -1;
4022 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004023
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004024 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4025 &element, &cnt))
4026 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004027
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004028 if (PyTuple_Size(args) == 2 && cnt < 0)
4029 cnt = 0;
4030
4031 ro = (repeatobject *)type->tp_alloc(type, 0);
4032 if (ro == NULL)
4033 return NULL;
4034 Py_INCREF(element);
4035 ro->element = element;
4036 ro->cnt = cnt;
4037 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004038}
4039
4040static void
4041repeat_dealloc(repeatobject *ro)
4042{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 PyObject_GC_UnTrack(ro);
4044 Py_XDECREF(ro->element);
4045 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004046}
4047
4048static int
4049repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4050{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004051 Py_VISIT(ro->element);
4052 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004053}
4054
4055static PyObject *
4056repeat_next(repeatobject *ro)
4057{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004058 if (ro->cnt == 0)
4059 return NULL;
4060 if (ro->cnt > 0)
4061 ro->cnt--;
4062 Py_INCREF(ro->element);
4063 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004064}
4065
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004066static PyObject *
4067repeat_repr(repeatobject *ro)
4068{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004069 if (ro->cnt == -1)
4070 return PyUnicode_FromFormat("repeat(%R)", ro->element);
4071 else
4072 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
4073}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004074
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004075static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004076repeat_len(repeatobject *ro)
4077{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004078 if (ro->cnt == -1) {
4079 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4080 return NULL;
4081 }
4082 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004083}
4084
Armin Rigof5b3e362006-02-11 21:32:43 +00004085PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004086
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004087static PyObject *
4088repeat_reduce(repeatobject *ro)
4089{
4090 /* unpickle this so that a new repeat iterator is constructed with an
4091 * object, then call __setstate__ on it to set cnt
4092 */
4093 if (ro->cnt >= 0)
4094 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4095 else
4096 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4097}
4098
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004099static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004100 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004101 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004102 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004103};
4104
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004105PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004106"repeat(object [,times]) -> create an iterator which returns the object\n\
4107for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004108endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004109
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004110static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004111 PyVarObject_HEAD_INIT(NULL, 0)
4112 "itertools.repeat", /* tp_name */
4113 sizeof(repeatobject), /* tp_basicsize */
4114 0, /* tp_itemsize */
4115 /* methods */
4116 (destructor)repeat_dealloc, /* tp_dealloc */
4117 0, /* tp_print */
4118 0, /* tp_getattr */
4119 0, /* tp_setattr */
4120 0, /* tp_reserved */
4121 (reprfunc)repeat_repr, /* tp_repr */
4122 0, /* tp_as_number */
4123 0, /* tp_as_sequence */
4124 0, /* tp_as_mapping */
4125 0, /* tp_hash */
4126 0, /* tp_call */
4127 0, /* tp_str */
4128 PyObject_GenericGetAttr, /* tp_getattro */
4129 0, /* tp_setattro */
4130 0, /* tp_as_buffer */
4131 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4132 Py_TPFLAGS_BASETYPE, /* tp_flags */
4133 repeat_doc, /* tp_doc */
4134 (traverseproc)repeat_traverse, /* tp_traverse */
4135 0, /* tp_clear */
4136 0, /* tp_richcompare */
4137 0, /* tp_weaklistoffset */
4138 PyObject_SelfIter, /* tp_iter */
4139 (iternextfunc)repeat_next, /* tp_iternext */
4140 repeat_methods, /* tp_methods */
4141 0, /* tp_members */
4142 0, /* tp_getset */
4143 0, /* tp_base */
4144 0, /* tp_dict */
4145 0, /* tp_descr_get */
4146 0, /* tp_descr_set */
4147 0, /* tp_dictoffset */
4148 0, /* tp_init */
4149 0, /* tp_alloc */
4150 repeat_new, /* tp_new */
4151 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004152};
4153
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004154/* ziplongest object ************************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004155
4156#include "Python.h"
4157
4158typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004159 PyObject_HEAD
4160 Py_ssize_t tuplesize;
4161 Py_ssize_t numactive;
4162 PyObject *ittuple; /* tuple of iterators */
4163 PyObject *result;
4164 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004165} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004166
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004167static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004168
4169static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004170zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004171{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004172 ziplongestobject *lz;
4173 Py_ssize_t i;
4174 PyObject *ittuple; /* tuple of iterators */
4175 PyObject *result;
4176 PyObject *fillvalue = Py_None;
4177 Py_ssize_t tuplesize = PySequence_Length(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004178
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004179 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
4180 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
4181 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
4182 PyErr_SetString(PyExc_TypeError,
4183 "zip_longest() got an unexpected keyword argument");
4184 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004185 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004186 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004187
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004188 /* args must be a tuple */
4189 assert(PyTuple_Check(args));
Thomas Wouterscf297e42007-02-23 15:07:44 +00004190
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004191 /* obtain iterators */
4192 ittuple = PyTuple_New(tuplesize);
4193 if (ittuple == NULL)
4194 return NULL;
4195 for (i=0; i < tuplesize; ++i) {
4196 PyObject *item = PyTuple_GET_ITEM(args, i);
4197 PyObject *it = PyObject_GetIter(item);
4198 if (it == NULL) {
4199 if (PyErr_ExceptionMatches(PyExc_TypeError))
4200 PyErr_Format(PyExc_TypeError,
4201 "zip_longest argument #%zd must support iteration",
4202 i+1);
4203 Py_DECREF(ittuple);
4204 return NULL;
4205 }
4206 PyTuple_SET_ITEM(ittuple, i, it);
4207 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004208
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004209 /* create a result holder */
4210 result = PyTuple_New(tuplesize);
4211 if (result == NULL) {
4212 Py_DECREF(ittuple);
4213 return NULL;
4214 }
4215 for (i=0 ; i < tuplesize ; i++) {
4216 Py_INCREF(Py_None);
4217 PyTuple_SET_ITEM(result, i, Py_None);
4218 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004219
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004220 /* create ziplongestobject structure */
4221 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4222 if (lz == NULL) {
4223 Py_DECREF(ittuple);
4224 Py_DECREF(result);
4225 return NULL;
4226 }
4227 lz->ittuple = ittuple;
4228 lz->tuplesize = tuplesize;
4229 lz->numactive = tuplesize;
4230 lz->result = result;
4231 Py_INCREF(fillvalue);
4232 lz->fillvalue = fillvalue;
4233 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004234}
4235
4236static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004237zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004238{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004239 PyObject_GC_UnTrack(lz);
4240 Py_XDECREF(lz->ittuple);
4241 Py_XDECREF(lz->result);
4242 Py_XDECREF(lz->fillvalue);
4243 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004244}
4245
4246static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004247zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004248{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004249 Py_VISIT(lz->ittuple);
4250 Py_VISIT(lz->result);
4251 Py_VISIT(lz->fillvalue);
4252 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004253}
4254
4255static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004256zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004257{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004258 Py_ssize_t i;
4259 Py_ssize_t tuplesize = lz->tuplesize;
4260 PyObject *result = lz->result;
4261 PyObject *it;
4262 PyObject *item;
4263 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004265 if (tuplesize == 0)
4266 return NULL;
4267 if (lz->numactive == 0)
4268 return NULL;
4269 if (Py_REFCNT(result) == 1) {
4270 Py_INCREF(result);
4271 for (i=0 ; i < tuplesize ; i++) {
4272 it = PyTuple_GET_ITEM(lz->ittuple, i);
4273 if (it == NULL) {
4274 Py_INCREF(lz->fillvalue);
4275 item = lz->fillvalue;
4276 } else {
4277 item = PyIter_Next(it);
4278 if (item == NULL) {
4279 lz->numactive -= 1;
4280 if (lz->numactive == 0 || PyErr_Occurred()) {
4281 lz->numactive = 0;
4282 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004283 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004284 } else {
4285 Py_INCREF(lz->fillvalue);
4286 item = lz->fillvalue;
4287 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4288 Py_DECREF(it);
4289 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004290 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004291 }
4292 olditem = PyTuple_GET_ITEM(result, i);
4293 PyTuple_SET_ITEM(result, i, item);
4294 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004295 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004296 } else {
4297 result = PyTuple_New(tuplesize);
4298 if (result == NULL)
4299 return NULL;
4300 for (i=0 ; i < tuplesize ; i++) {
4301 it = PyTuple_GET_ITEM(lz->ittuple, i);
4302 if (it == NULL) {
4303 Py_INCREF(lz->fillvalue);
4304 item = lz->fillvalue;
4305 } else {
4306 item = PyIter_Next(it);
4307 if (item == NULL) {
4308 lz->numactive -= 1;
4309 if (lz->numactive == 0 || PyErr_Occurred()) {
4310 lz->numactive = 0;
4311 Py_DECREF(result);
4312 return NULL;
4313 } else {
4314 Py_INCREF(lz->fillvalue);
4315 item = lz->fillvalue;
4316 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4317 Py_DECREF(it);
4318 }
4319 }
4320 }
4321 PyTuple_SET_ITEM(result, i, item);
4322 }
4323 }
4324 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004325}
4326
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004327static PyObject *
4328zip_longest_reduce(ziplongestobject *lz)
4329{
4330
4331 /* Create a new tuple with empty sequences where appropriate to pickle.
4332 * Then use setstate to set the fillvalue
4333 */
4334 int i;
4335 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4336 if (args == NULL)
4337 return NULL;
4338 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4339 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4340 if (elem == NULL) {
4341 elem = PyTuple_New(0);
4342 if (elem == NULL) {
4343 Py_DECREF(args);
4344 return NULL;
4345 }
4346 } else
4347 Py_INCREF(elem);
4348 PyTuple_SET_ITEM(args, i, elem);
4349 }
4350 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4351}
4352
4353static PyObject *
4354zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4355{
4356 Py_CLEAR(lz->fillvalue);
4357 lz->fillvalue = state;
4358 Py_INCREF(lz->fillvalue);
4359 Py_RETURN_NONE;
4360}
4361
4362static PyMethodDef zip_longest_methods[] = {
4363 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4364 reduce_doc},
4365 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4366 setstate_doc},
4367 {NULL, NULL} /* sentinel */
4368};
4369
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004370PyDoc_STRVAR(zip_longest_doc,
4371"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004372\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004373Return an zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004374the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004375method continues until the longest iterable in the argument sequence\n\
4376is exhausted and then it raises StopIteration. When the shorter iterables\n\
4377are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4378defaults to None or can be specified by a keyword argument.\n\
4379");
4380
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004381static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004382 PyVarObject_HEAD_INIT(NULL, 0)
4383 "itertools.zip_longest", /* tp_name */
4384 sizeof(ziplongestobject), /* tp_basicsize */
4385 0, /* tp_itemsize */
4386 /* methods */
4387 (destructor)zip_longest_dealloc, /* tp_dealloc */
4388 0, /* tp_print */
4389 0, /* tp_getattr */
4390 0, /* tp_setattr */
4391 0, /* tp_reserved */
4392 0, /* tp_repr */
4393 0, /* tp_as_number */
4394 0, /* tp_as_sequence */
4395 0, /* tp_as_mapping */
4396 0, /* tp_hash */
4397 0, /* tp_call */
4398 0, /* tp_str */
4399 PyObject_GenericGetAttr, /* tp_getattro */
4400 0, /* tp_setattro */
4401 0, /* tp_as_buffer */
4402 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4403 Py_TPFLAGS_BASETYPE, /* tp_flags */
4404 zip_longest_doc, /* tp_doc */
4405 (traverseproc)zip_longest_traverse, /* tp_traverse */
4406 0, /* tp_clear */
4407 0, /* tp_richcompare */
4408 0, /* tp_weaklistoffset */
4409 PyObject_SelfIter, /* tp_iter */
4410 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004411 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004412 0, /* tp_members */
4413 0, /* tp_getset */
4414 0, /* tp_base */
4415 0, /* tp_dict */
4416 0, /* tp_descr_get */
4417 0, /* tp_descr_set */
4418 0, /* tp_dictoffset */
4419 0, /* tp_init */
4420 0, /* tp_alloc */
4421 zip_longest_new, /* tp_new */
4422 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004423};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004424
4425/* module level code ********************************************************/
4426
4427PyDoc_STRVAR(module_doc,
4428"Functional tools for creating and using iterators.\n\
4429\n\
4430Infinite iterators:\n\
4431count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004432cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004433repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004434\n\
4435Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004436accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004437chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
4438compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4439dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4440groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004441filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004442islice(seq, [start,] stop [, step]) --> elements from\n\
4443 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004444starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004445tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004446takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004447zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004448\n\
4449Combinatoric generators:\n\
4450product(p, q, ... [repeat=1]) --> cartesian product\n\
4451permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004452combinations(p, r)\n\
4453combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004454");
4455
4456
Raymond Hettingerad983e72003-11-12 14:32:26 +00004457static PyMethodDef module_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004458 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4459 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004460};
4461
Martin v. Löwis1a214512008-06-11 05:26:20 +00004462
4463static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004464 PyModuleDef_HEAD_INIT,
4465 "itertools",
4466 module_doc,
4467 -1,
4468 module_methods,
4469 NULL,
4470 NULL,
4471 NULL,
4472 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004473};
4474
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004475PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004476PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004477{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004478 int i;
4479 PyObject *m;
4480 char *name;
4481 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004482 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004483 &combinations_type,
4484 &cwr_type,
4485 &cycle_type,
4486 &dropwhile_type,
4487 &takewhile_type,
4488 &islice_type,
4489 &starmap_type,
4490 &chain_type,
4491 &compress_type,
4492 &filterfalse_type,
4493 &count_type,
4494 &ziplongest_type,
4495 &permutations_type,
4496 &product_type,
4497 &repeat_type,
4498 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004499 &_grouper_type,
4500 &tee_type,
4501 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004502 NULL
4503 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004504
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004505 Py_TYPE(&teedataobject_type) = &PyType_Type;
4506 m = PyModule_Create(&itertoolsmodule);
4507 if (m == NULL)
4508 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004509
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004510 for (i=0 ; typelist[i] != NULL ; i++) {
4511 if (PyType_Ready(typelist[i]) < 0)
4512 return NULL;
4513 name = strchr(typelist[i]->tp_name, '.');
4514 assert (name != NULL);
4515 Py_INCREF(typelist[i]);
4516 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4517 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004518
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004519 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004520}