blob: 91b482d3bdfe186947a52780278252a25849162d [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 Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
Antoine Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +000098 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000101
Antoine Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +0000114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000117
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122
Antoine Pitrouc7c96a92010-05-09 15:15:40 +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 Pitrouc7c96a92010-05-09 15:15:40 +0000128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000131
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_compare */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000199 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000200
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000208
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000209 PyObject_GC_Track(igo);
210 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000228}
229
230static PyObject *
231_grouper_next(_grouperobject *igo)
232{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000236
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000241
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
251 }
252 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000253
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000269 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000270}
271
272static PyTypeObject _grouper_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_compare */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
294 0, /* tp_doc */
295 (traverseproc)_grouper_traverse,/* tp_traverse */
296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000315
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319/* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000321 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000330 PyObject_HEAD
331 PyObject *it;
332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000338 PyObject_HEAD
339 teedataobject *dataobj;
340 int index;
341 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000343
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000345
346static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000347teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000348{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
359 PyObject_GC_Track(tdo);
360 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
368 Py_XINCREF(tdo->nextlink);
369 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000375 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000376
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
388 }
389 Py_INCREF(value);
390 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391}
392
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000393static int
394teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000413}
414
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000415static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000416teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421}
422
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425static PyTypeObject teedataobject_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
429 0, /* tp_itemsize */
430 /* methods */
431 (destructor)teedataobject_dealloc, /* tp_dealloc */
432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_compare */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
447 teedataobject_doc, /* tp_doc */
448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466};
467
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000468
469static PyTypeObject tee_type;
470
471static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000472tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000473{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000474 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
478 Py_DECREF(to->dataobj);
479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487}
488
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000499 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
507 newto->weakreflist = NULL;
508 PyObject_GC_Track(newto);
509 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000517 teeobject *to;
518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
526 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000527
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000537
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000542 Py_XDECREF(it);
543 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
553 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554}
555
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000556static int
557tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
561 Py_CLEAR(to->dataobj);
562 return 0;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000563}
564
565static void
566tee_dealloc(teeobject *to)
567{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000571}
572
Raymond Hettingerad983e72003-11-12 14:32:26 +0000573PyDoc_STRVAR(teeobject_doc,
574"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
Raymond Hettingerad983e72003-11-12 14:32:26 +0000576static PyMethodDef tee_methods[] = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_compare */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
599 0, /* tp_getattro */
600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
603 teeobject_doc, /* tp_doc */
604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
606 0, /* tp_richcompare */
607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
621 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000622};
623
Raymond Hettingerad983e72003-11-12 14:32:26 +0000624static PyObject *
625tee(PyObject *self, PyObject *args)
626{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000629
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
631 return NULL;
632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
652 }
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000665}
666
667PyDoc_STRVAR(tee_doc,
668"tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671/* cycle object **********************************************************/
672
673typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000678} cycleobject;
679
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000680static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000681
682static PyObject *
683cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000692
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000700
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000706
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000717
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000718 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
727 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728}
729
730static int
731cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
735 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass)
749 PyList_Append(lz->saved, item);
750 return item;
751 }
752 if (PyErr_Occurred()) {
753 if (PyErr_ExceptionMatches(PyExc_StopIteration))
754 PyErr_Clear();
755 else
756 return NULL;
757 }
758 if (PyList_Size(lz->saved) == 0)
759 return NULL;
760 it = PyObject_GetIter(lz->saved);
761 if (it == NULL)
762 return NULL;
763 tmp = lz->it;
764 lz->it = it;
765 lz->firstpass = 1;
766 Py_DECREF(tmp);
767 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000768}
769
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000770PyDoc_STRVAR(cycle_doc,
771"cycle(iterable) --> cycle object\n\
772\n\
773Return elements from the iterable until it is exhausted.\n\
774Then repeat the sequence indefinitely.");
775
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000776static PyTypeObject cycle_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000777 PyVarObject_HEAD_INIT(NULL, 0)
778 "itertools.cycle", /* tp_name */
779 sizeof(cycleobject), /* tp_basicsize */
780 0, /* tp_itemsize */
781 /* methods */
782 (destructor)cycle_dealloc, /* tp_dealloc */
783 0, /* tp_print */
784 0, /* tp_getattr */
785 0, /* tp_setattr */
786 0, /* tp_compare */
787 0, /* tp_repr */
788 0, /* tp_as_number */
789 0, /* tp_as_sequence */
790 0, /* tp_as_mapping */
791 0, /* tp_hash */
792 0, /* tp_call */
793 0, /* tp_str */
794 PyObject_GenericGetAttr, /* tp_getattro */
795 0, /* tp_setattro */
796 0, /* tp_as_buffer */
797 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
798 Py_TPFLAGS_BASETYPE, /* tp_flags */
799 cycle_doc, /* tp_doc */
800 (traverseproc)cycle_traverse, /* tp_traverse */
801 0, /* tp_clear */
802 0, /* tp_richcompare */
803 0, /* tp_weaklistoffset */
804 PyObject_SelfIter, /* tp_iter */
805 (iternextfunc)cycle_next, /* tp_iternext */
806 0, /* tp_methods */
807 0, /* tp_members */
808 0, /* tp_getset */
809 0, /* tp_base */
810 0, /* tp_dict */
811 0, /* tp_descr_get */
812 0, /* tp_descr_set */
813 0, /* tp_dictoffset */
814 0, /* tp_init */
815 0, /* tp_alloc */
816 cycle_new, /* tp_new */
817 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000818};
819
820
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000821/* dropwhile object **********************************************************/
822
823typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000824 PyObject_HEAD
825 PyObject *func;
826 PyObject *it;
827 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000828} dropwhileobject;
829
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000830static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831
832static PyObject *
833dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
834{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000835 PyObject *func, *seq;
836 PyObject *it;
837 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000838
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000839 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
840 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000841
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000842 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
843 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000844
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000845 /* Get iterator. */
846 it = PyObject_GetIter(seq);
847 if (it == NULL)
848 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000849
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000850 /* create dropwhileobject structure */
851 lz = (dropwhileobject *)type->tp_alloc(type, 0);
852 if (lz == NULL) {
853 Py_DECREF(it);
854 return NULL;
855 }
856 Py_INCREF(func);
857 lz->func = func;
858 lz->it = it;
859 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000860
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000861 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862}
863
864static void
865dropwhile_dealloc(dropwhileobject *lz)
866{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000867 PyObject_GC_UnTrack(lz);
868 Py_XDECREF(lz->func);
869 Py_XDECREF(lz->it);
870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871}
872
873static int
874dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
875{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000876 Py_VISIT(lz->it);
877 Py_VISIT(lz->func);
878 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000879}
880
881static PyObject *
882dropwhile_next(dropwhileobject *lz)
883{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000884 PyObject *item, *good;
885 PyObject *it = lz->it;
886 long ok;
887 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000889 assert(PyIter_Check(it));
890 iternext = *Py_TYPE(it)->tp_iternext;
891 for (;;) {
892 item = iternext(it);
893 if (item == NULL)
894 return NULL;
895 if (lz->start == 1)
896 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000897
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000898 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
899 if (good == NULL) {
900 Py_DECREF(item);
901 return NULL;
902 }
903 ok = PyObject_IsTrue(good);
904 Py_DECREF(good);
905 if (!ok) {
906 lz->start = 1;
907 return item;
908 }
909 Py_DECREF(item);
910 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000911}
912
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000913PyDoc_STRVAR(dropwhile_doc,
914"dropwhile(predicate, iterable) --> dropwhile object\n\
915\n\
916Drop items from the iterable while predicate(item) is true.\n\
917Afterwards, return every element until the iterable is exhausted.");
918
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000919static PyTypeObject dropwhile_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000920 PyVarObject_HEAD_INIT(NULL, 0)
921 "itertools.dropwhile", /* tp_name */
922 sizeof(dropwhileobject), /* tp_basicsize */
923 0, /* tp_itemsize */
924 /* methods */
925 (destructor)dropwhile_dealloc, /* tp_dealloc */
926 0, /* tp_print */
927 0, /* tp_getattr */
928 0, /* tp_setattr */
929 0, /* tp_compare */
930 0, /* tp_repr */
931 0, /* tp_as_number */
932 0, /* tp_as_sequence */
933 0, /* tp_as_mapping */
934 0, /* tp_hash */
935 0, /* tp_call */
936 0, /* tp_str */
937 PyObject_GenericGetAttr, /* tp_getattro */
938 0, /* tp_setattro */
939 0, /* tp_as_buffer */
940 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
941 Py_TPFLAGS_BASETYPE, /* tp_flags */
942 dropwhile_doc, /* tp_doc */
943 (traverseproc)dropwhile_traverse, /* tp_traverse */
944 0, /* tp_clear */
945 0, /* tp_richcompare */
946 0, /* tp_weaklistoffset */
947 PyObject_SelfIter, /* tp_iter */
948 (iternextfunc)dropwhile_next, /* tp_iternext */
949 0, /* tp_methods */
950 0, /* tp_members */
951 0, /* tp_getset */
952 0, /* tp_base */
953 0, /* tp_dict */
954 0, /* tp_descr_get */
955 0, /* tp_descr_set */
956 0, /* tp_dictoffset */
957 0, /* tp_init */
958 0, /* tp_alloc */
959 dropwhile_new, /* tp_new */
960 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000961};
962
963
964/* takewhile object **********************************************************/
965
966typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000967 PyObject_HEAD
968 PyObject *func;
969 PyObject *it;
970 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000971} takewhileobject;
972
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000973static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000974
975static PyObject *
976takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
977{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000978 PyObject *func, *seq;
979 PyObject *it;
980 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000981
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000982 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
983 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +0000984
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000985 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
986 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000987
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000988 /* Get iterator. */
989 it = PyObject_GetIter(seq);
990 if (it == NULL)
991 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000992
Antoine Pitrouc7c96a92010-05-09 15:15:40 +0000993 /* create takewhileobject structure */
994 lz = (takewhileobject *)type->tp_alloc(type, 0);
995 if (lz == NULL) {
996 Py_DECREF(it);
997 return NULL;
998 }
999 Py_INCREF(func);
1000 lz->func = func;
1001 lz->it = it;
1002 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001003
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001004 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001005}
1006
1007static void
1008takewhile_dealloc(takewhileobject *lz)
1009{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001010 PyObject_GC_UnTrack(lz);
1011 Py_XDECREF(lz->func);
1012 Py_XDECREF(lz->it);
1013 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014}
1015
1016static int
1017takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1018{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001019 Py_VISIT(lz->it);
1020 Py_VISIT(lz->func);
1021 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001022}
1023
1024static PyObject *
1025takewhile_next(takewhileobject *lz)
1026{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001027 PyObject *item, *good;
1028 PyObject *it = lz->it;
1029 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001030
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001031 if (lz->stop == 1)
1032 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001033
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001034 assert(PyIter_Check(it));
1035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001051}
1052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053PyDoc_STRVAR(takewhile_doc,
1054"takewhile(predicate, iterable) --> takewhile object\n\
1055\n\
1056Return successive entries from an iterable as long as the \n\
1057predicate evaluates to true for each entry.");
1058
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001059static PyTypeObject takewhile_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
1061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
1063 0, /* tp_itemsize */
1064 /* methods */
1065 (destructor)takewhile_dealloc, /* tp_dealloc */
1066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_compare */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
1082 takewhile_doc, /* tp_doc */
1083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
1087 PyObject_SelfIter, /* tp_iter */
1088 (iternextfunc)takewhile_next, /* tp_iternext */
1089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
1098 0, /* tp_alloc */
1099 takewhile_new, /* tp_new */
1100 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001107 PyObject_HEAD
1108 PyObject *it;
1109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001113} isliceobject;
1114
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001115static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116
1117static PyObject *
1118islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001120 PyObject *seq;
1121 Py_ssize_t start=0, stop=-1, step=1;
1122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1123 Py_ssize_t numargs;
1124 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001128
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyInt_AsSsize_t(a1);
1136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
1140 "Stop argument for islice() must be a non-negative integer or None.");
1141 return NULL;
1142 }
1143 }
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyInt_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyInt_AsSsize_t(a2);
1151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
1155 "Stop argument for islice() must be a non-negative integer or None.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be non-negative integers or None.");
1163 return NULL;
1164 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyInt_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1176 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001195 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static int
1207islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001209 Py_VISIT(lz->it);
1210 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001216 PyObject *item;
1217 PyObject *it = lz->it;
1218 Py_ssize_t oldnext;
1219 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001221 assert(PyIter_Check(it));
1222 iternext = *Py_TYPE(it)->tp_iternext;
1223 while (lz->cnt < lz->next) {
1224 item = iternext(it);
1225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
1230 if (lz->stop != -1 && lz->cnt >= lz->stop)
1231 return NULL;
1232 assert(PyIter_Check(it));
1233 item = iternext(it);
1234 if (item == NULL)
1235 return NULL;
1236 lz->cnt++;
1237 oldnext = lz->next;
1238 lz->next += lz->step;
1239 if (lz->next < oldnext) /* Check for overflow */
1240 lz->next = lz->stop;
1241 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242}
1243
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244PyDoc_STRVAR(islice_doc,
1245"islice(iterable, [start,] stop [, step]) --> islice object\n\
1246\n\
1247Return an iterator whose next() method returns selected values from an\n\
1248iterable. If start is specified, will skip all preceding elements;\n\
1249otherwise, start defaults to zero. Step defaults to one. If\n\
1250specified as another value, step determines how many values are \n\
1251skipped between successive calls. Works like a slice() on a list\n\
1252but returns an iterator.");
1253
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001254static PyTypeObject islice_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001255 PyVarObject_HEAD_INIT(NULL, 0)
1256 "itertools.islice", /* tp_name */
1257 sizeof(isliceobject), /* tp_basicsize */
1258 0, /* tp_itemsize */
1259 /* methods */
1260 (destructor)islice_dealloc, /* tp_dealloc */
1261 0, /* tp_print */
1262 0, /* tp_getattr */
1263 0, /* tp_setattr */
1264 0, /* tp_compare */
1265 0, /* tp_repr */
1266 0, /* tp_as_number */
1267 0, /* tp_as_sequence */
1268 0, /* tp_as_mapping */
1269 0, /* tp_hash */
1270 0, /* tp_call */
1271 0, /* tp_str */
1272 PyObject_GenericGetAttr, /* tp_getattro */
1273 0, /* tp_setattro */
1274 0, /* tp_as_buffer */
1275 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1276 Py_TPFLAGS_BASETYPE, /* tp_flags */
1277 islice_doc, /* tp_doc */
1278 (traverseproc)islice_traverse, /* tp_traverse */
1279 0, /* tp_clear */
1280 0, /* tp_richcompare */
1281 0, /* tp_weaklistoffset */
1282 PyObject_SelfIter, /* tp_iter */
1283 (iternextfunc)islice_next, /* tp_iternext */
1284 0, /* tp_methods */
1285 0, /* tp_members */
1286 0, /* tp_getset */
1287 0, /* tp_base */
1288 0, /* tp_dict */
1289 0, /* tp_descr_get */
1290 0, /* tp_descr_set */
1291 0, /* tp_dictoffset */
1292 0, /* tp_init */
1293 0, /* tp_alloc */
1294 islice_new, /* tp_new */
1295 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001296};
1297
1298
1299/* starmap object ************************************************************/
1300
1301typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001302 PyObject_HEAD
1303 PyObject *func;
1304 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001305} starmapobject;
1306
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001307static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001308
1309static PyObject *
1310starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1311{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001312 PyObject *func, *seq;
1313 PyObject *it;
1314 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001315
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001316 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1317 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001318
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001319 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1320 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001322 /* Get iterator. */
1323 it = PyObject_GetIter(seq);
1324 if (it == NULL)
1325 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001327 /* create starmapobject structure */
1328 lz = (starmapobject *)type->tp_alloc(type, 0);
1329 if (lz == NULL) {
1330 Py_DECREF(it);
1331 return NULL;
1332 }
1333 Py_INCREF(func);
1334 lz->func = func;
1335 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001336
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001337 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001338}
1339
1340static void
1341starmap_dealloc(starmapobject *lz)
1342{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001343 PyObject_GC_UnTrack(lz);
1344 Py_XDECREF(lz->func);
1345 Py_XDECREF(lz->it);
1346 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001347}
1348
1349static int
1350starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1351{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001352 Py_VISIT(lz->it);
1353 Py_VISIT(lz->func);
1354 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001355}
1356
1357static PyObject *
1358starmap_next(starmapobject *lz)
1359{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001360 PyObject *args;
1361 PyObject *result;
1362 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001364 assert(PyIter_Check(it));
1365 args = (*Py_TYPE(it)->tp_iternext)(it);
1366 if (args == NULL)
1367 return NULL;
1368 if (!PyTuple_CheckExact(args)) {
1369 PyObject *newargs = PySequence_Tuple(args);
1370 Py_DECREF(args);
1371 if (newargs == NULL)
1372 return NULL;
1373 args = newargs;
1374 }
1375 result = PyObject_Call(lz->func, args, NULL);
1376 Py_DECREF(args);
1377 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378}
1379
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380PyDoc_STRVAR(starmap_doc,
1381"starmap(function, sequence) --> starmap object\n\
1382\n\
1383Return an iterator whose values are returned from the function evaluated\n\
1384with a argument tuple taken from the given sequence.");
1385
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001386static PyTypeObject starmap_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001387 PyVarObject_HEAD_INIT(NULL, 0)
1388 "itertools.starmap", /* tp_name */
1389 sizeof(starmapobject), /* tp_basicsize */
1390 0, /* tp_itemsize */
1391 /* methods */
1392 (destructor)starmap_dealloc, /* tp_dealloc */
1393 0, /* tp_print */
1394 0, /* tp_getattr */
1395 0, /* tp_setattr */
1396 0, /* tp_compare */
1397 0, /* tp_repr */
1398 0, /* tp_as_number */
1399 0, /* tp_as_sequence */
1400 0, /* tp_as_mapping */
1401 0, /* tp_hash */
1402 0, /* tp_call */
1403 0, /* tp_str */
1404 PyObject_GenericGetAttr, /* tp_getattro */
1405 0, /* tp_setattro */
1406 0, /* tp_as_buffer */
1407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1408 Py_TPFLAGS_BASETYPE, /* tp_flags */
1409 starmap_doc, /* tp_doc */
1410 (traverseproc)starmap_traverse, /* tp_traverse */
1411 0, /* tp_clear */
1412 0, /* tp_richcompare */
1413 0, /* tp_weaklistoffset */
1414 PyObject_SelfIter, /* tp_iter */
1415 (iternextfunc)starmap_next, /* tp_iternext */
1416 0, /* tp_methods */
1417 0, /* tp_members */
1418 0, /* tp_getset */
1419 0, /* tp_base */
1420 0, /* tp_dict */
1421 0, /* tp_descr_get */
1422 0, /* tp_descr_set */
1423 0, /* tp_dictoffset */
1424 0, /* tp_init */
1425 0, /* tp_alloc */
1426 starmap_new, /* tp_new */
1427 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001428};
1429
1430
1431/* imap object ************************************************************/
1432
1433typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001434 PyObject_HEAD
1435 PyObject *iters;
1436 PyObject *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437} imapobject;
1438
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001439static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440
1441static PyObject *
1442imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1443{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001444 PyObject *it, *iters, *func;
1445 imapobject *lz;
1446 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001448 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1449 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001450
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001451 numargs = PyTuple_Size(args);
1452 if (numargs < 2) {
1453 PyErr_SetString(PyExc_TypeError,
1454 "imap() must have at least two arguments.");
1455 return NULL;
1456 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001458 iters = PyTuple_New(numargs-1);
1459 if (iters == NULL)
1460 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001462 for (i=1 ; i<numargs ; i++) {
1463 /* Get iterator. */
1464 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1465 if (it == NULL) {
1466 Py_DECREF(iters);
1467 return NULL;
1468 }
1469 PyTuple_SET_ITEM(iters, i-1, it);
1470 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001472 /* create imapobject structure */
1473 lz = (imapobject *)type->tp_alloc(type, 0);
1474 if (lz == NULL) {
1475 Py_DECREF(iters);
1476 return NULL;
1477 }
1478 lz->iters = iters;
1479 func = PyTuple_GET_ITEM(args, 0);
1480 Py_INCREF(func);
1481 lz->func = func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001483 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static void
1487imap_dealloc(imapobject *lz)
1488{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001489 PyObject_GC_UnTrack(lz);
1490 Py_XDECREF(lz->iters);
1491 Py_XDECREF(lz->func);
1492 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001493}
1494
1495static int
1496imap_traverse(imapobject *lz, visitproc visit, void *arg)
1497{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001498 Py_VISIT(lz->iters);
1499 Py_VISIT(lz->func);
1500 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501}
1502
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001503/*
Raymond Hettinger2012f172003-02-07 05:32:58 +00001504imap() is an iterator version of __builtins__.map() except that it does
1505not have the None fill-in feature. That was intentionally left out for
1506the following reasons:
1507
1508 1) Itertools are designed to be easily combined and chained together.
1509 Having all tools stop with the shortest input is a unifying principle
1510 that makes it easier to combine finite iterators (supplying data) with
1511 infinite iterators like count() and repeat() (for supplying sequential
1512 or constant arguments to a function).
1513
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001514 2) In typical use cases for combining itertools, having one finite data
Raymond Hettinger2012f172003-02-07 05:32:58 +00001515 supplier run out before another is likely to be an error condition which
1516 should not pass silently by automatically supplying None.
1517
1518 3) The use cases for automatic None fill-in are rare -- not many functions
1519 do something useful when a parameter suddenly switches type and becomes
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001520 None.
Raymond Hettinger2012f172003-02-07 05:32:58 +00001521
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001522 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001523 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001524
1525 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1526*/
1527
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528static PyObject *
1529imap_next(imapobject *lz)
1530{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001531 PyObject *val;
1532 PyObject *argtuple;
1533 PyObject *result;
1534 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001536 numargs = PyTuple_Size(lz->iters);
1537 argtuple = PyTuple_New(numargs);
1538 if (argtuple == NULL)
1539 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001541 for (i=0 ; i<numargs ; i++) {
1542 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1543 if (val == NULL) {
1544 Py_DECREF(argtuple);
1545 return NULL;
1546 }
1547 PyTuple_SET_ITEM(argtuple, i, val);
1548 }
1549 if (lz->func == Py_None)
1550 return argtuple;
1551 result = PyObject_Call(lz->func, argtuple, NULL);
1552 Py_DECREF(argtuple);
1553 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554}
1555
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556PyDoc_STRVAR(imap_doc,
1557"imap(func, *iterables) --> imap object\n\
1558\n\
1559Make an iterator that computes the function using arguments from\n\
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001560each of the iterables. Like map() except that it returns\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001561an iterator instead of a list and that it stops when the shortest\n\
1562iterable is exhausted instead of filling in None for shorter\n\
1563iterables.");
1564
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001565static PyTypeObject imap_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001566 PyVarObject_HEAD_INIT(NULL, 0)
1567 "itertools.imap", /* tp_name */
1568 sizeof(imapobject), /* tp_basicsize */
1569 0, /* tp_itemsize */
1570 /* methods */
1571 (destructor)imap_dealloc, /* tp_dealloc */
1572 0, /* tp_print */
1573 0, /* tp_getattr */
1574 0, /* tp_setattr */
1575 0, /* tp_compare */
1576 0, /* tp_repr */
1577 0, /* tp_as_number */
1578 0, /* tp_as_sequence */
1579 0, /* tp_as_mapping */
1580 0, /* tp_hash */
1581 0, /* tp_call */
1582 0, /* tp_str */
1583 PyObject_GenericGetAttr, /* tp_getattro */
1584 0, /* tp_setattro */
1585 0, /* tp_as_buffer */
1586 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1587 Py_TPFLAGS_BASETYPE, /* tp_flags */
1588 imap_doc, /* tp_doc */
1589 (traverseproc)imap_traverse, /* tp_traverse */
1590 0, /* tp_clear */
1591 0, /* tp_richcompare */
1592 0, /* tp_weaklistoffset */
1593 PyObject_SelfIter, /* tp_iter */
1594 (iternextfunc)imap_next, /* tp_iternext */
1595 0, /* tp_methods */
1596 0, /* tp_members */
1597 0, /* tp_getset */
1598 0, /* tp_base */
1599 0, /* tp_dict */
1600 0, /* tp_descr_get */
1601 0, /* tp_descr_set */
1602 0, /* tp_dictoffset */
1603 0, /* tp_init */
1604 0, /* tp_alloc */
1605 imap_new, /* tp_new */
1606 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607};
1608
1609
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001610/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611
1612typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001613 PyObject_HEAD
1614 PyObject *source; /* Iterator over input iterables */
1615 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001616} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001618static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001620static PyObject *
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001621chain_new_internal(PyTypeObject *type, PyObject *source)
1622{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001623 chainobject *lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001624
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001625 lz = (chainobject *)type->tp_alloc(type, 0);
1626 if (lz == NULL) {
1627 Py_DECREF(source);
1628 return NULL;
1629 }
1630
1631 lz->source = source;
1632 lz->active = NULL;
1633 return (PyObject *)lz;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001634}
1635
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001637chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001639 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001641 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1642 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001644 source = PyObject_GetIter(args);
1645 if (source == NULL)
1646 return NULL;
1647
1648 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001649}
1650
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001651static PyObject *
1652chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1653{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001654 PyObject *source;
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001655
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001656 source = PyObject_GetIter(arg);
1657 if (source == NULL)
1658 return NULL;
1659
1660 return chain_new_internal(type, source);
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001661}
1662
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001663static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001664chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001666 PyObject_GC_UnTrack(lz);
1667 Py_XDECREF(lz->active);
1668 Py_XDECREF(lz->source);
1669 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670}
1671
Raymond Hettinger2012f172003-02-07 05:32:58 +00001672static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001673chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001674{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001675 Py_VISIT(lz->source);
1676 Py_VISIT(lz->active);
1677 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001678}
1679
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001681chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001683 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001685 if (lz->source == NULL)
1686 return NULL; /* already stopped */
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001687
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001688 if (lz->active == NULL) {
1689 PyObject *iterable = PyIter_Next(lz->source);
1690 if (iterable == NULL) {
1691 Py_CLEAR(lz->source);
1692 return NULL; /* no more input sources */
1693 }
1694 lz->active = PyObject_GetIter(iterable);
1695 Py_DECREF(iterable);
1696 if (lz->active == NULL) {
1697 Py_CLEAR(lz->source);
1698 return NULL; /* input not iterable */
1699 }
1700 }
1701 item = PyIter_Next(lz->active);
1702 if (item != NULL)
1703 return item;
1704 if (PyErr_Occurred()) {
1705 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1706 PyErr_Clear();
1707 else
1708 return NULL; /* input raised an exception */
1709 }
1710 Py_CLEAR(lz->active);
1711 return chain_next(lz); /* recurse and use next active */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712}
1713
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001714PyDoc_STRVAR(chain_doc,
1715"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001717Return a chain object whose .next() method returns elements from the\n\
1718first iterable until it is exhausted, then elements from the next\n\
1719iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001720
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001721PyDoc_STRVAR(chain_from_iterable_doc,
1722"chain.from_iterable(iterable) --> chain object\n\
1723\n\
1724Alternate chain() contructor taking a single iterable argument\n\
1725that evaluates lazily.");
1726
1727static PyMethodDef chain_methods[] = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001728 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1729 chain_from_iterable_doc},
1730 {NULL, NULL} /* sentinel */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001731};
1732
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001733static PyTypeObject chain_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001734 PyVarObject_HEAD_INIT(NULL, 0)
1735 "itertools.chain", /* tp_name */
1736 sizeof(chainobject), /* tp_basicsize */
1737 0, /* tp_itemsize */
1738 /* methods */
1739 (destructor)chain_dealloc, /* tp_dealloc */
1740 0, /* tp_print */
1741 0, /* tp_getattr */
1742 0, /* tp_setattr */
1743 0, /* tp_compare */
1744 0, /* tp_repr */
1745 0, /* tp_as_number */
1746 0, /* tp_as_sequence */
1747 0, /* tp_as_mapping */
1748 0, /* tp_hash */
1749 0, /* tp_call */
1750 0, /* tp_str */
1751 PyObject_GenericGetAttr, /* tp_getattro */
1752 0, /* tp_setattro */
1753 0, /* tp_as_buffer */
1754 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1755 Py_TPFLAGS_BASETYPE, /* tp_flags */
1756 chain_doc, /* tp_doc */
1757 (traverseproc)chain_traverse, /* tp_traverse */
1758 0, /* tp_clear */
1759 0, /* tp_richcompare */
1760 0, /* tp_weaklistoffset */
1761 PyObject_SelfIter, /* tp_iter */
1762 (iternextfunc)chain_next, /* tp_iternext */
1763 chain_methods, /* tp_methods */
1764 0, /* tp_members */
1765 0, /* tp_getset */
1766 0, /* tp_base */
1767 0, /* tp_dict */
1768 0, /* tp_descr_get */
1769 0, /* tp_descr_set */
1770 0, /* tp_dictoffset */
1771 0, /* tp_init */
1772 0, /* tp_alloc */
1773 chain_new, /* tp_new */
1774 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001775};
1776
1777
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001778/* product object ************************************************************/
1779
1780typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001781 PyObject_HEAD
1782 PyObject *pools; /* tuple of pool tuples */
1783 Py_ssize_t *indices; /* one index per pool */
1784 PyObject *result; /* most recently returned result tuple */
1785 int stopped; /* set to 1 when the product iterator is exhausted */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001786} productobject;
1787
1788static PyTypeObject product_type;
1789
1790static PyObject *
1791product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1792{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001793 productobject *lz;
1794 Py_ssize_t nargs, npools, repeat=1;
1795 PyObject *pools = NULL;
1796 Py_ssize_t *indices = NULL;
1797 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001798
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001799 if (kwds != NULL) {
1800 char *kwlist[] = {"repeat", 0};
1801 PyObject *tmpargs = PyTuple_New(0);
1802 if (tmpargs == NULL)
1803 return NULL;
1804 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1805 Py_DECREF(tmpargs);
1806 return NULL;
1807 }
1808 Py_DECREF(tmpargs);
1809 if (repeat < 0) {
1810 PyErr_SetString(PyExc_ValueError,
1811 "repeat argument cannot be negative");
1812 return NULL;
1813 }
1814 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001815
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001816 assert(PyTuple_Check(args));
1817 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1818 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001819
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001820 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1821 if (indices == NULL) {
1822 PyErr_NoMemory();
1823 goto error;
1824 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001825
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001826 pools = PyTuple_New(npools);
1827 if (pools == NULL)
1828 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001829
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001830 for (i=0; i < nargs ; ++i) {
1831 PyObject *item = PyTuple_GET_ITEM(args, i);
1832 PyObject *pool = PySequence_Tuple(item);
1833 if (pool == NULL)
1834 goto error;
1835 PyTuple_SET_ITEM(pools, i, pool);
1836 indices[i] = 0;
1837 }
1838 for ( ; i < npools; ++i) {
1839 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1840 Py_INCREF(pool);
1841 PyTuple_SET_ITEM(pools, i, pool);
1842 indices[i] = 0;
1843 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001844
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001845 /* create productobject structure */
1846 lz = (productobject *)type->tp_alloc(type, 0);
1847 if (lz == NULL)
1848 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001849
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001850 lz->pools = pools;
1851 lz->indices = indices;
1852 lz->result = NULL;
1853 lz->stopped = 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001854
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001855 return (PyObject *)lz;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001856
1857error:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001858 if (indices != NULL)
1859 PyMem_Free(indices);
1860 Py_XDECREF(pools);
1861 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001862}
1863
1864static void
1865product_dealloc(productobject *lz)
1866{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001867 PyObject_GC_UnTrack(lz);
1868 Py_XDECREF(lz->pools);
1869 Py_XDECREF(lz->result);
1870 PyMem_Free(lz->indices);
1871 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001872}
1873
1874static int
1875product_traverse(productobject *lz, visitproc visit, void *arg)
1876{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001877 Py_VISIT(lz->pools);
1878 Py_VISIT(lz->result);
1879 return 0;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001880}
1881
1882static PyObject *
1883product_next(productobject *lz)
1884{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001885 PyObject *pool;
1886 PyObject *elem;
1887 PyObject *oldelem;
1888 PyObject *pools = lz->pools;
1889 PyObject *result = lz->result;
1890 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1891 Py_ssize_t i;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001892
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001893 if (lz->stopped)
1894 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001895
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001896 if (result == NULL) {
1897 /* On the first pass, return an initial tuple filled with the
1898 first element from each pool. */
1899 result = PyTuple_New(npools);
1900 if (result == NULL)
1901 goto empty;
1902 lz->result = result;
1903 for (i=0; i < npools; i++) {
1904 pool = PyTuple_GET_ITEM(pools, i);
1905 if (PyTuple_GET_SIZE(pool) == 0)
1906 goto empty;
1907 elem = PyTuple_GET_ITEM(pool, 0);
1908 Py_INCREF(elem);
1909 PyTuple_SET_ITEM(result, i, elem);
1910 }
1911 } else {
1912 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001913
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001914 /* Copy the previous result tuple or re-use it if available */
1915 if (Py_REFCNT(result) > 1) {
1916 PyObject *old_result = result;
1917 result = PyTuple_New(npools);
1918 if (result == NULL)
1919 goto empty;
1920 lz->result = result;
1921 for (i=0; i < npools; i++) {
1922 elem = PyTuple_GET_ITEM(old_result, i);
1923 Py_INCREF(elem);
1924 PyTuple_SET_ITEM(result, i, elem);
1925 }
1926 Py_DECREF(old_result);
1927 }
1928 /* Now, we've got the only copy so we can update it in-place */
1929 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001930
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001931 /* Update the pool indices right-to-left. Only advance to the
1932 next pool when the previous one rolls-over */
1933 for (i=npools-1 ; i >= 0 ; i--) {
1934 pool = PyTuple_GET_ITEM(pools, i);
1935 indices[i]++;
1936 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1937 /* Roll-over and advance to next pool */
1938 indices[i] = 0;
1939 elem = PyTuple_GET_ITEM(pool, 0);
1940 Py_INCREF(elem);
1941 oldelem = PyTuple_GET_ITEM(result, i);
1942 PyTuple_SET_ITEM(result, i, elem);
1943 Py_DECREF(oldelem);
1944 } else {
1945 /* No rollover. Just increment and stop here. */
1946 elem = PyTuple_GET_ITEM(pool, indices[i]);
1947 Py_INCREF(elem);
1948 oldelem = PyTuple_GET_ITEM(result, i);
1949 PyTuple_SET_ITEM(result, i, elem);
1950 Py_DECREF(oldelem);
1951 break;
1952 }
1953 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001954
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001955 /* If i is negative, then the indices have all rolled-over
1956 and we're done. */
1957 if (i < 0)
1958 goto empty;
1959 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001960
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001961 Py_INCREF(result);
1962 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001963
1964empty:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001965 lz->stopped = 1;
1966 return NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001967}
1968
1969PyDoc_STRVAR(product_doc,
1970"product(*iterables) --> product object\n\
1971\n\
Raymond Hettinger73d79632008-02-23 02:20:41 +00001972Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001973For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1974The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1975cycle in a manner similar to an odometer (with the rightmost element changing\n\
1976on every iteration).\n\n\
1977product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1978product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1979
1980static PyTypeObject product_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00001981 PyVarObject_HEAD_INIT(NULL, 0)
1982 "itertools.product", /* tp_name */
1983 sizeof(productobject), /* tp_basicsize */
1984 0, /* tp_itemsize */
1985 /* methods */
1986 (destructor)product_dealloc, /* tp_dealloc */
1987 0, /* tp_print */
1988 0, /* tp_getattr */
1989 0, /* tp_setattr */
1990 0, /* tp_compare */
1991 0, /* tp_repr */
1992 0, /* tp_as_number */
1993 0, /* tp_as_sequence */
1994 0, /* tp_as_mapping */
1995 0, /* tp_hash */
1996 0, /* tp_call */
1997 0, /* tp_str */
1998 PyObject_GenericGetAttr, /* tp_getattro */
1999 0, /* tp_setattro */
2000 0, /* tp_as_buffer */
2001 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2002 Py_TPFLAGS_BASETYPE, /* tp_flags */
2003 product_doc, /* tp_doc */
2004 (traverseproc)product_traverse, /* tp_traverse */
2005 0, /* tp_clear */
2006 0, /* tp_richcompare */
2007 0, /* tp_weaklistoffset */
2008 PyObject_SelfIter, /* tp_iter */
2009 (iternextfunc)product_next, /* tp_iternext */
2010 0, /* tp_methods */
2011 0, /* tp_members */
2012 0, /* tp_getset */
2013 0, /* tp_base */
2014 0, /* tp_dict */
2015 0, /* tp_descr_get */
2016 0, /* tp_descr_set */
2017 0, /* tp_dictoffset */
2018 0, /* tp_init */
2019 0, /* tp_alloc */
2020 product_new, /* tp_new */
2021 PyObject_GC_Del, /* tp_free */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00002022};
2023
2024
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002025/* combinations object ************************************************************/
2026
2027typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002028 PyObject_HEAD
2029 PyObject *pool; /* input converted to a tuple */
2030 Py_ssize_t *indices; /* one index per result element */
2031 PyObject *result; /* most recently returned result tuple */
2032 Py_ssize_t r; /* size of result tuple */
2033 int stopped; /* set to 1 when the combinations iterator is exhausted */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002034} combinationsobject;
2035
2036static PyTypeObject combinations_type;
2037
2038static PyObject *
2039combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2040{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002041 combinationsobject *co;
2042 Py_ssize_t n;
2043 Py_ssize_t r;
2044 PyObject *pool = NULL;
2045 PyObject *iterable = NULL;
2046 Py_ssize_t *indices = NULL;
2047 Py_ssize_t i;
2048 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002049
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002050 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2051 &iterable, &r))
2052 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002053
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002054 pool = PySequence_Tuple(iterable);
2055 if (pool == NULL)
2056 goto error;
2057 n = PyTuple_GET_SIZE(pool);
2058 if (r < 0) {
2059 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2060 goto error;
2061 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002062
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002063 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2064 if (indices == NULL) {
2065 PyErr_NoMemory();
2066 goto error;
2067 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002068
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002069 for (i=0 ; i<r ; i++)
2070 indices[i] = i;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002071
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002072 /* create combinationsobject structure */
2073 co = (combinationsobject *)type->tp_alloc(type, 0);
2074 if (co == NULL)
2075 goto error;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002076
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002077 co->pool = pool;
2078 co->indices = indices;
2079 co->result = NULL;
2080 co->r = r;
2081 co->stopped = r > n ? 1 : 0;
2082
2083 return (PyObject *)co;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002084
2085error:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002086 if (indices != NULL)
2087 PyMem_Free(indices);
2088 Py_XDECREF(pool);
2089 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002090}
2091
2092static void
2093combinations_dealloc(combinationsobject *co)
2094{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002095 PyObject_GC_UnTrack(co);
2096 Py_XDECREF(co->pool);
2097 Py_XDECREF(co->result);
2098 PyMem_Free(co->indices);
2099 Py_TYPE(co)->tp_free(co);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002100}
2101
2102static int
2103combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2104{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002105 Py_VISIT(co->pool);
2106 Py_VISIT(co->result);
2107 return 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002108}
2109
2110static PyObject *
2111combinations_next(combinationsobject *co)
2112{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002113 PyObject *elem;
2114 PyObject *oldelem;
2115 PyObject *pool = co->pool;
2116 Py_ssize_t *indices = co->indices;
2117 PyObject *result = co->result;
2118 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2119 Py_ssize_t r = co->r;
2120 Py_ssize_t i, j, index;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002121
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002122 if (co->stopped)
2123 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002124
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002125 if (result == NULL) {
2126 /* On the first pass, initialize result tuple using the indices */
2127 result = PyTuple_New(r);
2128 if (result == NULL)
2129 goto empty;
2130 co->result = result;
2131 for (i=0; i<r ; i++) {
2132 index = indices[i];
2133 elem = PyTuple_GET_ITEM(pool, index);
2134 Py_INCREF(elem);
2135 PyTuple_SET_ITEM(result, i, elem);
2136 }
2137 } else {
2138 /* Copy the previous result tuple or re-use it if available */
2139 if (Py_REFCNT(result) > 1) {
2140 PyObject *old_result = result;
2141 result = PyTuple_New(r);
2142 if (result == NULL)
2143 goto empty;
2144 co->result = result;
2145 for (i=0; i<r ; i++) {
2146 elem = PyTuple_GET_ITEM(old_result, i);
2147 Py_INCREF(elem);
2148 PyTuple_SET_ITEM(result, i, elem);
2149 }
2150 Py_DECREF(old_result);
2151 }
2152 /* Now, we've got the only copy so we can update it in-place
2153 * CPython's empty tuple is a singleton and cached in
2154 * PyTuple's freelist.
2155 */
2156 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002157
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002158 /* Scan indices right-to-left until finding one that is not
2159 at its maximum (i + n - r). */
2160 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2161 ;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002162
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002163 /* If i is negative, then the indices are all at
2164 their maximum value and we're done. */
2165 if (i < 0)
2166 goto empty;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002167
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002168 /* Increment the current index which we know is not at its
2169 maximum. Then move back to the right setting each index
2170 to its lowest possible value (one higher than the index
2171 to its left -- this maintains the sort order invariant). */
2172 indices[i]++;
2173 for (j=i+1 ; j<r ; j++)
2174 indices[j] = indices[j-1] + 1;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002175
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002176 /* Update the result tuple for the new indices
2177 starting with i, the leftmost index that changed */
2178 for ( ; i<r ; i++) {
2179 index = indices[i];
2180 elem = PyTuple_GET_ITEM(pool, index);
2181 Py_INCREF(elem);
2182 oldelem = PyTuple_GET_ITEM(result, i);
2183 PyTuple_SET_ITEM(result, i, elem);
2184 Py_DECREF(oldelem);
2185 }
2186 }
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002187
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002188 Py_INCREF(result);
2189 return result;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002190
2191empty:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002192 co->stopped = 1;
2193 return NULL;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002194}
2195
2196PyDoc_STRVAR(combinations_doc,
Raymond Hettinger1aef4442009-11-19 01:26:23 +00002197"combinations(iterable, r) --> combinations object\n\
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002198\n\
2199Return successive r-length combinations of elements in the iterable.\n\n\
2200combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2201
2202static PyTypeObject combinations_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002203 PyVarObject_HEAD_INIT(NULL, 0)
2204 "itertools.combinations", /* tp_name */
2205 sizeof(combinationsobject), /* tp_basicsize */
2206 0, /* tp_itemsize */
2207 /* methods */
2208 (destructor)combinations_dealloc, /* tp_dealloc */
2209 0, /* tp_print */
2210 0, /* tp_getattr */
2211 0, /* tp_setattr */
2212 0, /* tp_compare */
2213 0, /* tp_repr */
2214 0, /* tp_as_number */
2215 0, /* tp_as_sequence */
2216 0, /* tp_as_mapping */
2217 0, /* tp_hash */
2218 0, /* tp_call */
2219 0, /* tp_str */
2220 PyObject_GenericGetAttr, /* tp_getattro */
2221 0, /* tp_setattro */
2222 0, /* tp_as_buffer */
2223 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2224 Py_TPFLAGS_BASETYPE, /* tp_flags */
2225 combinations_doc, /* tp_doc */
2226 (traverseproc)combinations_traverse, /* tp_traverse */
2227 0, /* tp_clear */
2228 0, /* tp_richcompare */
2229 0, /* tp_weaklistoffset */
2230 PyObject_SelfIter, /* tp_iter */
2231 (iternextfunc)combinations_next, /* tp_iternext */
2232 0, /* tp_methods */
2233 0, /* tp_members */
2234 0, /* tp_getset */
2235 0, /* tp_base */
2236 0, /* tp_dict */
2237 0, /* tp_descr_get */
2238 0, /* tp_descr_set */
2239 0, /* tp_dictoffset */
2240 0, /* tp_init */
2241 0, /* tp_alloc */
2242 combinations_new, /* tp_new */
2243 PyObject_GC_Del, /* tp_free */
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002244};
2245
2246
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002247/* permutations object ************************************************************
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002248
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002249def permutations(iterable, r=None):
2250 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2251 pool = tuple(iterable)
2252 n = len(pool)
2253 r = n if r is None else r
2254 indices = range(n)
2255 cycles = range(n-r+1, n+1)[::-1]
2256 yield tuple(pool[i] for i in indices[:r])
2257 while n:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002258 for i in reversed(range(r)):
2259 cycles[i] -= 1
2260 if cycles[i] == 0:
2261 indices[i:] = indices[i+1:] + indices[i:i+1]
2262 cycles[i] = n - i
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002263 else:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002264 j = cycles[i]
2265 indices[i], indices[-j] = indices[-j], indices[i]
2266 yield tuple(pool[i] for i in indices[:r])
2267 break
2268 else:
2269 return
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002270*/
2271
2272typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002273 PyObject_HEAD
2274 PyObject *pool; /* input converted to a tuple */
2275 Py_ssize_t *indices; /* one index per element in the pool */
2276 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2277 PyObject *result; /* most recently returned result tuple */
2278 Py_ssize_t r; /* size of result tuple */
2279 int stopped; /* set to 1 when the permutations iterator is exhausted */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002280} permutationsobject;
2281
2282static PyTypeObject permutations_type;
2283
2284static PyObject *
2285permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2286{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002287 permutationsobject *po;
2288 Py_ssize_t n;
2289 Py_ssize_t r;
2290 PyObject *robj = Py_None;
2291 PyObject *pool = NULL;
2292 PyObject *iterable = NULL;
2293 Py_ssize_t *indices = NULL;
2294 Py_ssize_t *cycles = NULL;
2295 Py_ssize_t i;
2296 static char *kwargs[] = {"iterable", "r", NULL};
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002297
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002298 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2299 &iterable, &robj))
2300 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002301
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002302 pool = PySequence_Tuple(iterable);
2303 if (pool == NULL)
2304 goto error;
2305 n = PyTuple_GET_SIZE(pool);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002306
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002307 r = n;
2308 if (robj != Py_None) {
2309 r = PyInt_AsSsize_t(robj);
2310 if (r == -1 && PyErr_Occurred())
2311 goto error;
2312 }
2313 if (r < 0) {
2314 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2315 goto error;
2316 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002317
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002318 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2319 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2320 if (indices == NULL || cycles == NULL) {
2321 PyErr_NoMemory();
2322 goto error;
2323 }
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002324
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002325 for (i=0 ; i<n ; i++)
2326 indices[i] = i;
2327 for (i=0 ; i<r ; i++)
2328 cycles[i] = n - i;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002329
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002330 /* create permutationsobject structure */
2331 po = (permutationsobject *)type->tp_alloc(type, 0);
2332 if (po == NULL)
2333 goto error;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002334
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002335 po->pool = pool;
2336 po->indices = indices;
2337 po->cycles = cycles;
2338 po->result = NULL;
2339 po->r = r;
2340 po->stopped = r > n ? 1 : 0;
2341
2342 return (PyObject *)po;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002343
2344error:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002345 if (indices != NULL)
2346 PyMem_Free(indices);
2347 if (cycles != NULL)
2348 PyMem_Free(cycles);
2349 Py_XDECREF(pool);
2350 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002351}
2352
2353static void
2354permutations_dealloc(permutationsobject *po)
2355{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002356 PyObject_GC_UnTrack(po);
2357 Py_XDECREF(po->pool);
2358 Py_XDECREF(po->result);
2359 PyMem_Free(po->indices);
2360 PyMem_Free(po->cycles);
2361 Py_TYPE(po)->tp_free(po);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002362}
2363
2364static int
2365permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2366{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002367 Py_VISIT(po->pool);
2368 Py_VISIT(po->result);
2369 return 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002370}
2371
2372static PyObject *
2373permutations_next(permutationsobject *po)
2374{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002375 PyObject *elem;
2376 PyObject *oldelem;
2377 PyObject *pool = po->pool;
2378 Py_ssize_t *indices = po->indices;
2379 Py_ssize_t *cycles = po->cycles;
2380 PyObject *result = po->result;
2381 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2382 Py_ssize_t r = po->r;
2383 Py_ssize_t i, j, k, index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002384
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002385 if (po->stopped)
2386 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002387
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002388 if (result == NULL) {
2389 /* On the first pass, initialize result tuple using the indices */
2390 result = PyTuple_New(r);
2391 if (result == NULL)
2392 goto empty;
2393 po->result = result;
2394 for (i=0; i<r ; i++) {
2395 index = indices[i];
2396 elem = PyTuple_GET_ITEM(pool, index);
2397 Py_INCREF(elem);
2398 PyTuple_SET_ITEM(result, i, elem);
2399 }
2400 } else {
2401 if (n == 0)
2402 goto empty;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002403
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002404 /* Copy the previous result tuple or re-use it if available */
2405 if (Py_REFCNT(result) > 1) {
2406 PyObject *old_result = result;
2407 result = PyTuple_New(r);
2408 if (result == NULL)
2409 goto empty;
2410 po->result = result;
2411 for (i=0; i<r ; i++) {
2412 elem = PyTuple_GET_ITEM(old_result, i);
2413 Py_INCREF(elem);
2414 PyTuple_SET_ITEM(result, i, elem);
2415 }
2416 Py_DECREF(old_result);
2417 }
2418 /* Now, we've got the only copy so we can update it in-place */
2419 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002420
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002421 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2422 for (i=r-1 ; i>=0 ; i--) {
2423 cycles[i] -= 1;
2424 if (cycles[i] == 0) {
2425 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2426 index = indices[i];
2427 for (j=i ; j<n-1 ; j++)
2428 indices[j] = indices[j+1];
2429 indices[n-1] = index;
2430 cycles[i] = n - i;
2431 } else {
2432 j = cycles[i];
2433 index = indices[i];
2434 indices[i] = indices[n-j];
2435 indices[n-j] = index;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002436
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002437 for (k=i; k<r ; k++) {
2438 /* start with i, the leftmost element that changed */
2439 /* yield tuple(pool[k] for k in indices[:r]) */
2440 index = indices[k];
2441 elem = PyTuple_GET_ITEM(pool, index);
2442 Py_INCREF(elem);
2443 oldelem = PyTuple_GET_ITEM(result, k);
2444 PyTuple_SET_ITEM(result, k, elem);
2445 Py_DECREF(oldelem);
2446 }
2447 break;
2448 }
2449 }
2450 /* If i is negative, then the cycles have all
2451 rolled-over and we're done. */
2452 if (i < 0)
2453 goto empty;
2454 }
2455 Py_INCREF(result);
2456 return result;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002457
2458empty:
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002459 po->stopped = 1;
2460 return NULL;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002461}
2462
2463PyDoc_STRVAR(permutations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +00002464"permutations(iterable[, r]) --> permutations object\n\
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002465\n\
2466Return successive r-length permutations of elements in the iterable.\n\n\
Georg Brandl8550d692008-06-10 12:46:39 +00002467permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002468
2469static PyTypeObject permutations_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002470 PyVarObject_HEAD_INIT(NULL, 0)
2471 "itertools.permutations", /* tp_name */
2472 sizeof(permutationsobject), /* tp_basicsize */
2473 0, /* tp_itemsize */
2474 /* methods */
2475 (destructor)permutations_dealloc, /* tp_dealloc */
2476 0, /* tp_print */
2477 0, /* tp_getattr */
2478 0, /* tp_setattr */
2479 0, /* tp_compare */
2480 0, /* tp_repr */
2481 0, /* tp_as_number */
2482 0, /* tp_as_sequence */
2483 0, /* tp_as_mapping */
2484 0, /* tp_hash */
2485 0, /* tp_call */
2486 0, /* tp_str */
2487 PyObject_GenericGetAttr, /* tp_getattro */
2488 0, /* tp_setattro */
2489 0, /* tp_as_buffer */
2490 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2491 Py_TPFLAGS_BASETYPE, /* tp_flags */
2492 permutations_doc, /* tp_doc */
2493 (traverseproc)permutations_traverse, /* tp_traverse */
2494 0, /* tp_clear */
2495 0, /* tp_richcompare */
2496 0, /* tp_weaklistoffset */
2497 PyObject_SelfIter, /* tp_iter */
2498 (iternextfunc)permutations_next, /* tp_iternext */
2499 0, /* tp_methods */
2500 0, /* tp_members */
2501 0, /* tp_getset */
2502 0, /* tp_base */
2503 0, /* tp_dict */
2504 0, /* tp_descr_get */
2505 0, /* tp_descr_set */
2506 0, /* tp_dictoffset */
2507 0, /* tp_init */
2508 0, /* tp_alloc */
2509 permutations_new, /* tp_new */
2510 PyObject_GC_Del, /* tp_free */
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002511};
2512
2513
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002514/* ifilter object ************************************************************/
2515
2516typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002517 PyObject_HEAD
2518 PyObject *func;
2519 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002520} ifilterobject;
2521
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002522static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002523
2524static PyObject *
2525ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2526{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002527 PyObject *func, *seq;
2528 PyObject *it;
2529 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002530
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002531 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2532 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002533
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002534 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2535 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002536
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002537 /* Get iterator. */
2538 it = PyObject_GetIter(seq);
2539 if (it == NULL)
2540 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002541
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002542 /* create ifilterobject structure */
2543 lz = (ifilterobject *)type->tp_alloc(type, 0);
2544 if (lz == NULL) {
2545 Py_DECREF(it);
2546 return NULL;
2547 }
2548 Py_INCREF(func);
2549 lz->func = func;
2550 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002551
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002552 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002553}
2554
2555static void
2556ifilter_dealloc(ifilterobject *lz)
2557{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002558 PyObject_GC_UnTrack(lz);
2559 Py_XDECREF(lz->func);
2560 Py_XDECREF(lz->it);
2561 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002562}
2563
2564static int
2565ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2566{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002567 Py_VISIT(lz->it);
2568 Py_VISIT(lz->func);
2569 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002570}
2571
2572static PyObject *
2573ifilter_next(ifilterobject *lz)
2574{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002575 PyObject *item;
2576 PyObject *it = lz->it;
2577 long ok;
2578 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002579
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002580 assert(PyIter_Check(it));
2581 iternext = *Py_TYPE(it)->tp_iternext;
2582 for (;;) {
2583 item = iternext(it);
2584 if (item == NULL)
2585 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002586
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002587 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2588 ok = PyObject_IsTrue(item);
2589 } else {
2590 PyObject *good;
2591 good = PyObject_CallFunctionObjArgs(lz->func,
2592 item, NULL);
2593 if (good == NULL) {
2594 Py_DECREF(item);
2595 return NULL;
2596 }
2597 ok = PyObject_IsTrue(good);
2598 Py_DECREF(good);
2599 }
2600 if (ok)
2601 return item;
2602 Py_DECREF(item);
2603 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002604}
2605
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002606PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002607"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002608\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002609Return those items of sequence for which function(item) is true.\n\
2610If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002611
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002612static PyTypeObject ifilter_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002613 PyVarObject_HEAD_INIT(NULL, 0)
2614 "itertools.ifilter", /* tp_name */
2615 sizeof(ifilterobject), /* tp_basicsize */
2616 0, /* tp_itemsize */
2617 /* methods */
2618 (destructor)ifilter_dealloc, /* tp_dealloc */
2619 0, /* tp_print */
2620 0, /* tp_getattr */
2621 0, /* tp_setattr */
2622 0, /* tp_compare */
2623 0, /* tp_repr */
2624 0, /* tp_as_number */
2625 0, /* tp_as_sequence */
2626 0, /* tp_as_mapping */
2627 0, /* tp_hash */
2628 0, /* tp_call */
2629 0, /* tp_str */
2630 PyObject_GenericGetAttr, /* tp_getattro */
2631 0, /* tp_setattro */
2632 0, /* tp_as_buffer */
2633 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2634 Py_TPFLAGS_BASETYPE, /* tp_flags */
2635 ifilter_doc, /* tp_doc */
2636 (traverseproc)ifilter_traverse, /* tp_traverse */
2637 0, /* tp_clear */
2638 0, /* tp_richcompare */
2639 0, /* tp_weaklistoffset */
2640 PyObject_SelfIter, /* tp_iter */
2641 (iternextfunc)ifilter_next, /* tp_iternext */
2642 0, /* tp_methods */
2643 0, /* tp_members */
2644 0, /* tp_getset */
2645 0, /* tp_base */
2646 0, /* tp_dict */
2647 0, /* tp_descr_get */
2648 0, /* tp_descr_set */
2649 0, /* tp_dictoffset */
2650 0, /* tp_init */
2651 0, /* tp_alloc */
2652 ifilter_new, /* tp_new */
2653 PyObject_GC_Del, /* tp_free */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002654};
2655
2656
2657/* ifilterfalse object ************************************************************/
2658
2659typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002660 PyObject_HEAD
2661 PyObject *func;
2662 PyObject *it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002663} ifilterfalseobject;
2664
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002665static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002666
2667static PyObject *
2668ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2669{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002670 PyObject *func, *seq;
2671 PyObject *it;
2672 ifilterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002673
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002674 if (type == &ifilterfalse_type &&
2675 !_PyArg_NoKeywords("ifilterfalse()", kwds))
2676 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002677
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002678 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2679 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002680
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002681 /* Get iterator. */
2682 it = PyObject_GetIter(seq);
2683 if (it == NULL)
2684 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002685
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002686 /* create ifilterfalseobject structure */
2687 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
2688 if (lz == NULL) {
2689 Py_DECREF(it);
2690 return NULL;
2691 }
2692 Py_INCREF(func);
2693 lz->func = func;
2694 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002695
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002696 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002697}
2698
2699static void
2700ifilterfalse_dealloc(ifilterfalseobject *lz)
2701{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002702 PyObject_GC_UnTrack(lz);
2703 Py_XDECREF(lz->func);
2704 Py_XDECREF(lz->it);
2705 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002706}
2707
2708static int
2709ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
2710{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002711 Py_VISIT(lz->it);
2712 Py_VISIT(lz->func);
2713 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002714}
2715
2716static PyObject *
2717ifilterfalse_next(ifilterfalseobject *lz)
2718{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002719 PyObject *item;
2720 PyObject *it = lz->it;
2721 long ok;
2722 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002723
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002724 assert(PyIter_Check(it));
2725 iternext = *Py_TYPE(it)->tp_iternext;
2726 for (;;) {
2727 item = iternext(it);
2728 if (item == NULL)
2729 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002730
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002731 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2732 ok = PyObject_IsTrue(item);
2733 } else {
2734 PyObject *good;
2735 good = PyObject_CallFunctionObjArgs(lz->func,
2736 item, NULL);
2737 if (good == NULL) {
2738 Py_DECREF(item);
2739 return NULL;
2740 }
2741 ok = PyObject_IsTrue(good);
2742 Py_DECREF(good);
2743 }
2744 if (!ok)
2745 return item;
2746 Py_DECREF(item);
2747 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002748}
2749
Raymond Hettinger60eca932003-02-09 06:40:58 +00002750PyDoc_STRVAR(ifilterfalse_doc,
2751"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2752\n\
2753Return those items of sequence for which function(item) is false.\n\
2754If function is None, return the items that are false.");
2755
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002756static PyTypeObject ifilterfalse_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002757 PyVarObject_HEAD_INIT(NULL, 0)
2758 "itertools.ifilterfalse", /* tp_name */
2759 sizeof(ifilterfalseobject), /* tp_basicsize */
2760 0, /* tp_itemsize */
2761 /* methods */
2762 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2763 0, /* tp_print */
2764 0, /* tp_getattr */
2765 0, /* tp_setattr */
2766 0, /* tp_compare */
2767 0, /* tp_repr */
2768 0, /* tp_as_number */
2769 0, /* tp_as_sequence */
2770 0, /* tp_as_mapping */
2771 0, /* tp_hash */
2772 0, /* tp_call */
2773 0, /* tp_str */
2774 PyObject_GenericGetAttr, /* tp_getattro */
2775 0, /* tp_setattro */
2776 0, /* tp_as_buffer */
2777 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2778 Py_TPFLAGS_BASETYPE, /* tp_flags */
2779 ifilterfalse_doc, /* tp_doc */
2780 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2781 0, /* tp_clear */
2782 0, /* tp_richcompare */
2783 0, /* tp_weaklistoffset */
2784 PyObject_SelfIter, /* tp_iter */
2785 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2786 0, /* tp_methods */
2787 0, /* tp_members */
2788 0, /* tp_getset */
2789 0, /* tp_base */
2790 0, /* tp_dict */
2791 0, /* tp_descr_get */
2792 0, /* tp_descr_set */
2793 0, /* tp_dictoffset */
2794 0, /* tp_init */
2795 0, /* tp_alloc */
2796 ifilterfalse_new, /* tp_new */
2797 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002798};
2799
2800
2801/* count object ************************************************************/
2802
2803typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002804 PyObject_HEAD
2805 Py_ssize_t cnt;
2806 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002807} countobject;
2808
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002809static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002810
2811static PyObject *
2812count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2813{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002814 countobject *lz;
2815 Py_ssize_t cnt = 0;
2816 PyObject *cnt_arg = NULL;
2817 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002818
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002819 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
2820 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002821
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002822 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
2823 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002824
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002825 if (cnt_arg != NULL) {
2826 cnt = PyInt_AsSsize_t(cnt_arg);
2827 if (cnt == -1 && PyErr_Occurred()) {
2828 PyErr_Clear();
2829 if (!PyLong_Check(cnt_arg)) {
2830 PyErr_SetString(PyExc_TypeError, "an integer is required");
2831 return NULL;
2832 }
2833 long_cnt = cnt_arg;
2834 Py_INCREF(long_cnt);
2835 cnt = PY_SSIZE_T_MAX;
2836 }
2837 }
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002838
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002839 /* create countobject structure */
2840 lz = (countobject *)PyObject_New(countobject, &count_type);
2841 if (lz == NULL) {
2842 Py_XDECREF(long_cnt);
2843 return NULL;
2844 }
2845 lz->cnt = cnt;
2846 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002847
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002848 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002849}
2850
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002851static void
2852count_dealloc(countobject *lz)
2853{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002854 Py_XDECREF(lz->long_cnt);
2855 PyObject_Del(lz);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002856}
2857
2858static PyObject *
2859count_nextlong(countobject *lz)
2860{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002861 static PyObject *one = NULL;
2862 PyObject *cnt;
2863 PyObject *stepped_up;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002864
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002865 if (lz->long_cnt == NULL) {
2866 lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
2867 if (lz->long_cnt == NULL)
2868 return NULL;
2869 }
2870 if (one == NULL) {
2871 one = PyInt_FromLong(1);
2872 if (one == NULL)
2873 return NULL;
2874 }
2875 cnt = lz->long_cnt;
2876 assert(cnt != NULL);
2877 stepped_up = PyNumber_Add(cnt, one);
2878 if (stepped_up == NULL)
2879 return NULL;
2880 lz->long_cnt = stepped_up;
2881 return cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002882}
2883
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002884static PyObject *
2885count_next(countobject *lz)
2886{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002887 if (lz->cnt == PY_SSIZE_T_MAX)
2888 return count_nextlong(lz);
2889 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002890}
2891
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002892static PyObject *
2893count_repr(countobject *lz)
2894{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002895 PyObject *cnt_repr;
2896 PyObject *result;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002897
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002898 if (lz->cnt != PY_SSIZE_T_MAX)
2899 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002900
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002901 cnt_repr = PyObject_Repr(lz->long_cnt);
2902 if (cnt_repr == NULL)
2903 return NULL;
2904 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
2905 Py_DECREF(cnt_repr);
2906 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002907}
2908
Raymond Hettinger673827c2009-11-30 11:15:28 +00002909static PyObject *
2910count_reduce(countobject *lz)
2911{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002912 if (lz->cnt == PY_SSIZE_T_MAX)
2913 return Py_BuildValue("O(O)", Py_TYPE(lz), lz->long_cnt);
2914 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettinger673827c2009-11-30 11:15:28 +00002915}
2916
2917PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
2918
2919static PyMethodDef count_methods[] = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002920 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
2921 count_reduce_doc},
2922 {NULL, NULL} /* sentinel */
Raymond Hettinger673827c2009-11-30 11:15:28 +00002923};
2924
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002925PyDoc_STRVAR(count_doc,
2926"count([firstval]) --> count object\n\
2927\n\
2928Return a count object whose .next() method returns consecutive\n\
2929integers starting from zero or, if specified, from firstval.");
2930
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002931static PyTypeObject count_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002932 PyVarObject_HEAD_INIT(NULL, 0)
2933 "itertools.count", /* tp_name */
2934 sizeof(countobject), /* tp_basicsize */
2935 0, /* tp_itemsize */
2936 /* methods */
2937 (destructor)count_dealloc, /* tp_dealloc */
2938 0, /* tp_print */
2939 0, /* tp_getattr */
2940 0, /* tp_setattr */
2941 0, /* tp_compare */
2942 (reprfunc)count_repr, /* tp_repr */
2943 0, /* tp_as_number */
2944 0, /* tp_as_sequence */
2945 0, /* tp_as_mapping */
2946 0, /* tp_hash */
2947 0, /* tp_call */
2948 0, /* tp_str */
2949 PyObject_GenericGetAttr, /* tp_getattro */
2950 0, /* tp_setattro */
2951 0, /* tp_as_buffer */
2952 Py_TPFLAGS_DEFAULT, /* tp_flags */
2953 count_doc, /* tp_doc */
2954 0, /* tp_traverse */
2955 0, /* tp_clear */
2956 0, /* tp_richcompare */
2957 0, /* tp_weaklistoffset */
2958 PyObject_SelfIter, /* tp_iter */
2959 (iternextfunc)count_next, /* tp_iternext */
2960 count_methods, /* tp_methods */
2961 0, /* tp_members */
2962 0, /* tp_getset */
2963 0, /* tp_base */
2964 0, /* tp_dict */
2965 0, /* tp_descr_get */
2966 0, /* tp_descr_set */
2967 0, /* tp_dictoffset */
2968 0, /* tp_init */
2969 0, /* tp_alloc */
2970 count_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002971};
2972
2973
2974/* izip object ************************************************************/
2975
2976#include "Python.h"
2977
2978typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002979 PyObject_HEAD
2980 Py_ssize_t tuplesize;
2981 PyObject *ittuple; /* tuple of iterators */
2982 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002983} izipobject;
2984
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002985static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002986
2987static PyObject *
2988izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2989{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002990 izipobject *lz;
2991 Py_ssize_t i;
2992 PyObject *ittuple; /* tuple of iterators */
2993 PyObject *result;
2994 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002995
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002996 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
2997 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00002998
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00002999 /* args must be a tuple */
3000 assert(PyTuple_Check(args));
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003001
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003002 /* obtain iterators */
3003 ittuple = PyTuple_New(tuplesize);
3004 if (ittuple == NULL)
3005 return NULL;
3006 for (i=0; i < tuplesize; ++i) {
3007 PyObject *item = PyTuple_GET_ITEM(args, i);
3008 PyObject *it = PyObject_GetIter(item);
3009 if (it == NULL) {
3010 if (PyErr_ExceptionMatches(PyExc_TypeError))
3011 PyErr_Format(PyExc_TypeError,
3012 "izip argument #%zd must support iteration",
3013 i+1);
3014 Py_DECREF(ittuple);
3015 return NULL;
3016 }
3017 PyTuple_SET_ITEM(ittuple, i, it);
3018 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003019
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003020 /* create a result holder */
3021 result = PyTuple_New(tuplesize);
3022 if (result == NULL) {
3023 Py_DECREF(ittuple);
3024 return NULL;
3025 }
3026 for (i=0 ; i < tuplesize ; i++) {
3027 Py_INCREF(Py_None);
3028 PyTuple_SET_ITEM(result, i, Py_None);
3029 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003030
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003031 /* create izipobject structure */
3032 lz = (izipobject *)type->tp_alloc(type, 0);
3033 if (lz == NULL) {
3034 Py_DECREF(ittuple);
3035 Py_DECREF(result);
3036 return NULL;
3037 }
3038 lz->ittuple = ittuple;
3039 lz->tuplesize = tuplesize;
3040 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003041
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003042 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003043}
3044
3045static void
3046izip_dealloc(izipobject *lz)
3047{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003048 PyObject_GC_UnTrack(lz);
3049 Py_XDECREF(lz->ittuple);
3050 Py_XDECREF(lz->result);
3051 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003052}
3053
3054static int
3055izip_traverse(izipobject *lz, visitproc visit, void *arg)
3056{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003057 Py_VISIT(lz->ittuple);
3058 Py_VISIT(lz->result);
3059 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003060}
3061
3062static PyObject *
3063izip_next(izipobject *lz)
3064{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003065 Py_ssize_t i;
3066 Py_ssize_t tuplesize = lz->tuplesize;
3067 PyObject *result = lz->result;
3068 PyObject *it;
3069 PyObject *item;
3070 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003071
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003072 if (tuplesize == 0)
3073 return NULL;
3074 if (Py_REFCNT(result) == 1) {
3075 Py_INCREF(result);
3076 for (i=0 ; i < tuplesize ; i++) {
3077 it = PyTuple_GET_ITEM(lz->ittuple, i);
3078 assert(PyIter_Check(it));
3079 item = (*Py_TYPE(it)->tp_iternext)(it);
3080 if (item == NULL) {
3081 Py_DECREF(result);
3082 return NULL;
3083 }
3084 olditem = PyTuple_GET_ITEM(result, i);
3085 PyTuple_SET_ITEM(result, i, item);
3086 Py_DECREF(olditem);
3087 }
3088 } else {
3089 result = PyTuple_New(tuplesize);
3090 if (result == NULL)
3091 return NULL;
3092 for (i=0 ; i < tuplesize ; i++) {
3093 it = PyTuple_GET_ITEM(lz->ittuple, i);
3094 assert(PyIter_Check(it));
3095 item = (*Py_TYPE(it)->tp_iternext)(it);
3096 if (item == NULL) {
3097 Py_DECREF(result);
3098 return NULL;
3099 }
3100 PyTuple_SET_ITEM(result, i, item);
3101 }
3102 }
3103 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003104}
3105
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003106PyDoc_STRVAR(izip_doc,
3107"izip(iter1 [,iter2 [...]]) --> izip object\n\
3108\n\
3109Return a izip object whose .next() method returns a tuple where\n\
3110the i-th element comes from the i-th iterable argument. The .next()\n\
3111method continues until the shortest iterable in the argument sequence\n\
3112is exhausted and then it raises StopIteration. Works like the zip()\n\
3113function but consumes less memory by returning an iterator instead of\n\
3114a list.");
3115
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003116static PyTypeObject izip_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003117 PyVarObject_HEAD_INIT(NULL, 0)
3118 "itertools.izip", /* tp_name */
3119 sizeof(izipobject), /* tp_basicsize */
3120 0, /* tp_itemsize */
3121 /* methods */
3122 (destructor)izip_dealloc, /* tp_dealloc */
3123 0, /* tp_print */
3124 0, /* tp_getattr */
3125 0, /* tp_setattr */
3126 0, /* tp_compare */
3127 0, /* tp_repr */
3128 0, /* tp_as_number */
3129 0, /* tp_as_sequence */
3130 0, /* tp_as_mapping */
3131 0, /* tp_hash */
3132 0, /* tp_call */
3133 0, /* tp_str */
3134 PyObject_GenericGetAttr, /* tp_getattro */
3135 0, /* tp_setattro */
3136 0, /* tp_as_buffer */
3137 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3138 Py_TPFLAGS_BASETYPE, /* tp_flags */
3139 izip_doc, /* tp_doc */
3140 (traverseproc)izip_traverse, /* tp_traverse */
3141 0, /* tp_clear */
3142 0, /* tp_richcompare */
3143 0, /* tp_weaklistoffset */
3144 PyObject_SelfIter, /* tp_iter */
3145 (iternextfunc)izip_next, /* tp_iternext */
3146 0, /* tp_methods */
3147 0, /* tp_members */
3148 0, /* tp_getset */
3149 0, /* tp_base */
3150 0, /* tp_dict */
3151 0, /* tp_descr_get */
3152 0, /* tp_descr_set */
3153 0, /* tp_dictoffset */
3154 0, /* tp_init */
3155 0, /* tp_alloc */
3156 izip_new, /* tp_new */
3157 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003158};
3159
3160
3161/* repeat object ************************************************************/
3162
3163typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003164 PyObject_HEAD
3165 PyObject *element;
3166 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003167} repeatobject;
3168
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003169static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003170
3171static PyObject *
3172repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3173{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003174 repeatobject *ro;
3175 PyObject *element;
3176 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003177
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003178 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
3179 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00003180
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003181 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
3182 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003183
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003184 if (PyTuple_Size(args) == 2 && cnt < 0)
3185 cnt = 0;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003186
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003187 ro = (repeatobject *)type->tp_alloc(type, 0);
3188 if (ro == NULL)
3189 return NULL;
3190 Py_INCREF(element);
3191 ro->element = element;
3192 ro->cnt = cnt;
3193 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003194}
3195
3196static void
3197repeat_dealloc(repeatobject *ro)
3198{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003199 PyObject_GC_UnTrack(ro);
3200 Py_XDECREF(ro->element);
3201 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003202}
3203
3204static int
3205repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3206{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003207 Py_VISIT(ro->element);
3208 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003209}
3210
3211static PyObject *
3212repeat_next(repeatobject *ro)
3213{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003214 if (ro->cnt == 0)
3215 return NULL;
3216 if (ro->cnt > 0)
3217 ro->cnt--;
3218 Py_INCREF(ro->element);
3219 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003220}
3221
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003222static PyObject *
3223repeat_repr(repeatobject *ro)
3224{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003225 PyObject *result, *objrepr;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003226
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003227 objrepr = PyObject_Repr(ro->element);
3228 if (objrepr == NULL)
3229 return NULL;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003230
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003231 if (ro->cnt == -1)
3232 result = PyString_FromFormat("repeat(%s)",
3233 PyString_AS_STRING(objrepr));
3234 else
3235 result = PyString_FromFormat("repeat(%s, %zd)",
3236 PyString_AS_STRING(objrepr), ro->cnt);
3237 Py_DECREF(objrepr);
3238 return result;
3239}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003240
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003241static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003242repeat_len(repeatobject *ro)
3243{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003244 if (ro->cnt == -1) {
3245 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3246 return NULL;
3247 }
3248 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003249}
3250
Armin Rigof5b3e362006-02-11 21:32:43 +00003251PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003252
3253static PyMethodDef repeat_methods[] = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003254 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3255 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003256};
3257
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003258PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003259"repeat(element [,times]) -> create an iterator which returns the element\n\
3260for the specified number of times. If not specified, returns the element\n\
3261endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003262
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003263static PyTypeObject repeat_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003264 PyVarObject_HEAD_INIT(NULL, 0)
3265 "itertools.repeat", /* tp_name */
3266 sizeof(repeatobject), /* tp_basicsize */
3267 0, /* tp_itemsize */
3268 /* methods */
3269 (destructor)repeat_dealloc, /* tp_dealloc */
3270 0, /* tp_print */
3271 0, /* tp_getattr */
3272 0, /* tp_setattr */
3273 0, /* tp_compare */
3274 (reprfunc)repeat_repr, /* tp_repr */
3275 0, /* tp_as_number */
3276 0, /* tp_as_sequence */
3277 0, /* tp_as_mapping */
3278 0, /* tp_hash */
3279 0, /* tp_call */
3280 0, /* tp_str */
3281 PyObject_GenericGetAttr, /* tp_getattro */
3282 0, /* tp_setattro */
3283 0, /* tp_as_buffer */
3284 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3285 Py_TPFLAGS_BASETYPE, /* tp_flags */
3286 repeat_doc, /* tp_doc */
3287 (traverseproc)repeat_traverse, /* tp_traverse */
3288 0, /* tp_clear */
3289 0, /* tp_richcompare */
3290 0, /* tp_weaklistoffset */
3291 PyObject_SelfIter, /* tp_iter */
3292 (iternextfunc)repeat_next, /* tp_iternext */
3293 repeat_methods, /* tp_methods */
3294 0, /* tp_members */
3295 0, /* tp_getset */
3296 0, /* tp_base */
3297 0, /* tp_dict */
3298 0, /* tp_descr_get */
3299 0, /* tp_descr_set */
3300 0, /* tp_dictoffset */
3301 0, /* tp_init */
3302 0, /* tp_alloc */
3303 repeat_new, /* tp_new */
3304 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003305};
3306
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003307/* iziplongest object ************************************************************/
3308
3309#include "Python.h"
3310
3311typedef struct {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003312 PyObject_HEAD
3313 Py_ssize_t tuplesize;
3314 Py_ssize_t numactive;
3315 PyObject *ittuple; /* tuple of iterators */
3316 PyObject *result;
3317 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003318} iziplongestobject;
3319
3320static PyTypeObject iziplongest_type;
3321
3322static PyObject *
3323izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3324{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003325 iziplongestobject *lz;
3326 Py_ssize_t i;
3327 PyObject *ittuple; /* tuple of iterators */
3328 PyObject *result;
3329 PyObject *fillvalue = Py_None;
3330 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003331
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003332 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3333 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3334 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3335 PyErr_SetString(PyExc_TypeError,
3336 "izip_longest() got an unexpected keyword argument");
3337 return NULL;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003338 }
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003339 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003340
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003341 /* args must be a tuple */
3342 assert(PyTuple_Check(args));
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003343
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003344 /* obtain iterators */
3345 ittuple = PyTuple_New(tuplesize);
3346 if (ittuple == NULL)
3347 return NULL;
3348 for (i=0; i < tuplesize; ++i) {
3349 PyObject *item = PyTuple_GET_ITEM(args, i);
3350 PyObject *it = PyObject_GetIter(item);
3351 if (it == NULL) {
3352 if (PyErr_ExceptionMatches(PyExc_TypeError))
3353 PyErr_Format(PyExc_TypeError,
3354 "izip_longest argument #%zd must support iteration",
3355 i+1);
3356 Py_DECREF(ittuple);
3357 return NULL;
3358 }
3359 PyTuple_SET_ITEM(ittuple, i, it);
3360 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003361
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003362 /* create a result holder */
3363 result = PyTuple_New(tuplesize);
3364 if (result == NULL) {
3365 Py_DECREF(ittuple);
3366 return NULL;
3367 }
3368 for (i=0 ; i < tuplesize ; i++) {
3369 Py_INCREF(Py_None);
3370 PyTuple_SET_ITEM(result, i, Py_None);
3371 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003372
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003373 /* create iziplongestobject structure */
3374 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3375 if (lz == NULL) {
3376 Py_DECREF(ittuple);
3377 Py_DECREF(result);
3378 return NULL;
3379 }
3380 lz->ittuple = ittuple;
3381 lz->tuplesize = tuplesize;
3382 lz->numactive = tuplesize;
3383 lz->result = result;
3384 Py_INCREF(fillvalue);
3385 lz->fillvalue = fillvalue;
3386 return (PyObject *)lz;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003387}
3388
3389static void
3390izip_longest_dealloc(iziplongestobject *lz)
3391{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003392 PyObject_GC_UnTrack(lz);
3393 Py_XDECREF(lz->ittuple);
3394 Py_XDECREF(lz->result);
3395 Py_XDECREF(lz->fillvalue);
3396 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003397}
3398
3399static int
3400izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3401{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003402 Py_VISIT(lz->ittuple);
3403 Py_VISIT(lz->result);
3404 Py_VISIT(lz->fillvalue);
3405 return 0;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003406}
3407
3408static PyObject *
3409izip_longest_next(iziplongestobject *lz)
3410{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003411 Py_ssize_t i;
3412 Py_ssize_t tuplesize = lz->tuplesize;
3413 PyObject *result = lz->result;
3414 PyObject *it;
3415 PyObject *item;
3416 PyObject *olditem;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003417
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003418 if (tuplesize == 0)
3419 return NULL;
3420 if (lz->numactive == 0)
3421 return NULL;
3422 if (Py_REFCNT(result) == 1) {
3423 Py_INCREF(result);
3424 for (i=0 ; i < tuplesize ; i++) {
3425 it = PyTuple_GET_ITEM(lz->ittuple, i);
3426 if (it == NULL) {
3427 Py_INCREF(lz->fillvalue);
3428 item = lz->fillvalue;
3429 } else {
3430 assert(PyIter_Check(it));
3431 item = PyIter_Next(it);
3432 if (item == NULL) {
3433 lz->numactive -= 1;
3434 if (lz->numactive == 0 || PyErr_Occurred()) {
3435 lz->numactive = 0;
3436 Py_DECREF(result);
Raymond Hettinger4da5faa2009-11-01 18:43:31 +00003437 return NULL;
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003438 } else {
3439 Py_INCREF(lz->fillvalue);
3440 item = lz->fillvalue;
3441 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3442 Py_DECREF(it);
3443 }
Raymond Hettinger4da5faa2009-11-01 18:43:31 +00003444 }
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003445 }
3446 olditem = PyTuple_GET_ITEM(result, i);
3447 PyTuple_SET_ITEM(result, i, item);
3448 Py_DECREF(olditem);
Raymond Hettinger743d8312009-11-01 20:05:41 +00003449 }
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003450 } else {
3451 result = PyTuple_New(tuplesize);
3452 if (result == NULL)
3453 return NULL;
3454 for (i=0 ; i < tuplesize ; i++) {
3455 it = PyTuple_GET_ITEM(lz->ittuple, i);
3456 if (it == NULL) {
3457 Py_INCREF(lz->fillvalue);
3458 item = lz->fillvalue;
3459 } else {
3460 assert(PyIter_Check(it));
3461 item = PyIter_Next(it);
3462 if (item == NULL) {
3463 lz->numactive -= 1;
3464 if (lz->numactive == 0 || PyErr_Occurred()) {
3465 lz->numactive = 0;
3466 Py_DECREF(result);
3467 return NULL;
3468 } else {
3469 Py_INCREF(lz->fillvalue);
3470 item = lz->fillvalue;
3471 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3472 Py_DECREF(it);
3473 }
3474 }
3475 }
3476 PyTuple_SET_ITEM(result, i, item);
3477 }
3478 }
3479 return result;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003480}
3481
3482PyDoc_STRVAR(izip_longest_doc,
3483"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3484\n\
3485Return an izip_longest object whose .next() method returns a tuple where\n\
3486the i-th element comes from the i-th iterable argument. The .next()\n\
3487method continues until the longest iterable in the argument sequence\n\
3488is exhausted and then it raises StopIteration. When the shorter iterables\n\
3489are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3490defaults to None or can be specified by a keyword argument.\n\
3491");
3492
3493static PyTypeObject iziplongest_type = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003494 PyVarObject_HEAD_INIT(NULL, 0)
3495 "itertools.izip_longest", /* tp_name */
3496 sizeof(iziplongestobject), /* tp_basicsize */
3497 0, /* tp_itemsize */
3498 /* methods */
3499 (destructor)izip_longest_dealloc, /* tp_dealloc */
3500 0, /* tp_print */
3501 0, /* tp_getattr */
3502 0, /* tp_setattr */
3503 0, /* tp_compare */
3504 0, /* tp_repr */
3505 0, /* tp_as_number */
3506 0, /* tp_as_sequence */
3507 0, /* tp_as_mapping */
3508 0, /* tp_hash */
3509 0, /* tp_call */
3510 0, /* tp_str */
3511 PyObject_GenericGetAttr, /* tp_getattro */
3512 0, /* tp_setattro */
3513 0, /* tp_as_buffer */
3514 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3515 Py_TPFLAGS_BASETYPE, /* tp_flags */
3516 izip_longest_doc, /* tp_doc */
3517 (traverseproc)izip_longest_traverse, /* tp_traverse */
3518 0, /* tp_clear */
3519 0, /* tp_richcompare */
3520 0, /* tp_weaklistoffset */
3521 PyObject_SelfIter, /* tp_iter */
3522 (iternextfunc)izip_longest_next, /* tp_iternext */
3523 0, /* tp_methods */
3524 0, /* tp_members */
3525 0, /* tp_getset */
3526 0, /* tp_base */
3527 0, /* tp_dict */
3528 0, /* tp_descr_get */
3529 0, /* tp_descr_set */
3530 0, /* tp_dictoffset */
3531 0, /* tp_init */
3532 0, /* tp_alloc */
3533 izip_longest_new, /* tp_new */
3534 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003535};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003536
3537/* module level code ********************************************************/
3538
3539PyDoc_STRVAR(module_doc,
3540"Functional tools for creating and using iterators.\n\
3541\n\
3542Infinite iterators:\n\
3543count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003544cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003545repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003546\n\
3547Iterators terminating on the shortest input sequence:\n\
3548izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003549izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003550ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3551ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003552islice(seq, [start,] stop [, step]) --> elements from\n\
3553 seq[start:stop:step]\n\
3554imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3555starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003556tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003557chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003558takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3559dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003560groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger1aef4442009-11-19 01:26:23 +00003561\n\
3562Combinatoric generators:\n\
3563product(p, q, ... [repeat=1]) --> cartesian product\n\
3564permutations(p[, r])\n\
3565combinations(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003566");
3567
3568
Raymond Hettingerad983e72003-11-12 14:32:26 +00003569static PyMethodDef module_methods[] = {
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003570 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3571 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00003572};
3573
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003574PyMODINIT_FUNC
3575inititertools(void)
3576{
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003577 int i;
3578 PyObject *m;
3579 char *name;
3580 PyTypeObject *typelist[] = {
3581 &combinations_type,
3582 &cycle_type,
3583 &dropwhile_type,
3584 &takewhile_type,
3585 &islice_type,
3586 &starmap_type,
3587 &imap_type,
3588 &chain_type,
3589 &ifilter_type,
3590 &ifilterfalse_type,
3591 &count_type,
3592 &izip_type,
3593 &iziplongest_type,
3594 &permutations_type,
3595 &product_type,
3596 &repeat_type,
3597 &groupby_type,
3598 NULL
3599 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00003600
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003601 Py_TYPE(&teedataobject_type) = &PyType_Type;
3602 m = Py_InitModule3("itertools", module_methods, module_doc);
3603 if (m == NULL)
3604 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003605
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003606 for (i=0 ; typelist[i] != NULL ; i++) {
3607 if (PyType_Ready(typelist[i]) < 0)
3608 return;
3609 name = strchr(typelist[i]->tp_name, '.');
3610 assert (name != NULL);
3611 Py_INCREF(typelist[i]);
3612 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3613 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003614
Antoine Pitrouc7c96a92010-05-09 15:15:40 +00003615 if (PyType_Ready(&teedataobject_type) < 0)
3616 return;
3617 if (PyType_Ready(&tee_type) < 0)
3618 return;
3619 if (PyType_Ready(&_grouper_type) < 0)
3620 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003621}