blob: 5875d1045800602a245cedde662f7a2ae6dc52b6 [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
5/* Itertools module written and maintained
6 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 {
15 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
21} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Martin v. Löwis02cbf4a2006-02-27 17:20:04 +000029 static char *kwargs[] = {"iterable", "key", NULL};
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
32
33 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;
51}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
56 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);
Christian Heimese93237d2007-12-19 02:37:44 +000062 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +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);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000073 return 0;
74}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
81 /* 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;
89
90 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 }
97
98 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
101
102 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 }
113
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000114 tmp = gbo->currkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000115 gbo->currkey = newkey;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000116 Py_XDECREF(tmp);
117
118 tmp = gbo->currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000119 gbo->currvalue = newvalue;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000120 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121 }
122
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000123 Py_INCREF(gbo->currkey);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
131
132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
135}
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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000143 "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 */
183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
199 _grouperobject *igo;
200
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
208
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000209 PyObject_GC_Track(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000210 return (PyObject *)igo;
211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000216 PyObject_GC_UnTrack(igo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000219 PyObject_GC_Del(igo);
220}
221
222static int
223_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224{
225 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{
233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
236
237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
241
242 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 }
253
254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
258
259 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;
264
265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000267 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000268
269 return r;
270}
271
272static PyTypeObject _grouper_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000273 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000274 "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 */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000294 0, /* tp_doc */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000295 (traverseproc)_grouper_traverse,/* tp_traverse */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000296 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 */
Raymond Hettingera1ca94a2008-03-06 22:51:36 +0000312 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000313};
314
315
316
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
321 pointers), the value should be a multiple of 16 minus space for
322 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 {
330 PyObject_HEAD
331 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
335} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
337typedef struct {
338 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339 teedataobject *dataobj;
340 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000341 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{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000349 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000350
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000352 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000353 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000354
355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000359 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000360 return (PyObject *)tdo;
361}
362
363static PyObject *
364teedataobject_jumplink(teedataobject *tdo)
365{
366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000368 Py_XINCREF(tdo->nextlink);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000369 return tdo->nextlink;
370}
371
372static PyObject *
373teedataobject_getitem(teedataobject *tdo, int i)
374{
375 PyObject *value;
376
377 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;
Raymond Hettinger45143692003-10-25 06:37:47 +0000388 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 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{
396 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;
402}
403
404static int
405teedataobject_clear(teedataobject *tdo)
406{
407 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;
413}
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{
Thomas Wouters19bf33b2006-03-27 21:02:13 +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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000429 0, /* tp_itemsize */
430 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432 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 */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000447 teedataobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000448 (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{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000474 PyObject *value, *link;
475
476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000478 Py_DECREF(to->dataobj);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000479 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;
487}
488
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000489static int
490tee_traverse(teeobject *to, visitproc visit, void *arg)
491{
492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
494}
495
Raymond Hettingerad983e72003-11-12 14:32:26 +0000496static PyObject *
497tee_copy(teeobject *to)
498{
499 teeobject *newto;
500
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000501 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000507 newto->weakreflist = NULL;
Thomas Woutersb3deb942006-04-15 22:33:13 +0000508 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000509 return (PyObject *)newto;
510}
511
512PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514static PyObject *
515tee_fromiterable(PyObject *iterable)
516{
517 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000518 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519
520 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 }
527
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000528 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
536 }
537
Raymond Hettingerad983e72003-11-12 14:32:26 +0000538 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000539 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000540 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000541done:
542 Py_XDECREF(it);
543 return (PyObject *)to;
544}
545
546static PyObject *
547tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000549 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000550
551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000553 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{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000561 Py_CLEAR(to->dataobj);
562 return 0;
563}
564
565static void
566tee_dealloc(teeobject *to)
567{
568 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[] = {
577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
579};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000580
581static PyTypeObject tee_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000582 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000583 "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 */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000599 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000603 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000606 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000611 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 */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000621 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{
Neal Norwitz69e88972006-09-02 02:58:13 +0000627 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000628 PyObject *it, *iterable, *copyable, *result;
629
Neal Norwitz69e88972006-09-02 02:58:13 +0000630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000631 return NULL;
Neal Norwitz69e88972006-09-02 02:58:13 +0000632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
635 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000636 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;
665}
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 {
674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
678} 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{
685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
689
Georg Brandlb84c1372007-01-21 10:28:43 +0000690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000691 return NULL;
692
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
695
696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
700
701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
706
707 /* 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;
717
718 return (PyObject *)lz;
719}
720
721static void
722cycle_dealloc(cycleobject *lz)
723{
724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000727 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000735 return 0;
736}
737
738static PyObject *
739cycle_next(cycleobject *lz)
740{
741 PyObject *item;
742 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000743 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000744
745 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 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000752 if (PyErr_Occurred()) {
753 if (PyErr_ExceptionMatches(PyExc_StopIteration))
754 PyErr_Clear();
755 else
756 return NULL;
757 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000758 if (PyList_Size(lz->saved) == 0)
759 return NULL;
760 it = PyObject_GetIter(lz->saved);
761 if (it == NULL)
762 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000763 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000764 lz->it = it;
765 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000766 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000767 }
768}
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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000777 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000778 "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 */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000804 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000805 (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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000815 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000816 cycle_new, /* tp_new */
817 PyObject_GC_Del, /* tp_free */
818};
819
820
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000821/* dropwhile object **********************************************************/
822
823typedef struct {
824 PyObject_HEAD
825 PyObject *func;
826 PyObject *it;
827 long start;
828} 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{
835 PyObject *func, *seq;
836 PyObject *it;
837 dropwhileobject *lz;
838
Georg Brandlb84c1372007-01-21 10:28:43 +0000839 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000840 return NULL;
841
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
843 return NULL;
844
845 /* Get iterator. */
846 it = PyObject_GetIter(seq);
847 if (it == NULL)
848 return NULL;
849
850 /* 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;
860
861 return (PyObject *)lz;
862}
863
864static void
865dropwhile_dealloc(dropwhileobject *lz)
866{
867 PyObject_GC_UnTrack(lz);
868 Py_XDECREF(lz->func);
869 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +0000870 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000876 Py_VISIT(lz->it);
877 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878 return 0;
879}
880
881static PyObject *
882dropwhile_next(dropwhileobject *lz)
883{
884 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000885 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000886 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000887 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000888
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000889 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +0000890 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000891 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000892 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000893 if (item == NULL)
894 return NULL;
895 if (lz->start == 1)
896 return item;
897
898 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 }
911}
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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000920 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000921 "itertools.dropwhile", /* tp_name */
922 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000923 0, /* tp_itemsize */
924 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000925 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000926 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 */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000942 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000943 (traverseproc)dropwhile_traverse, /* tp_traverse */
944 0, /* tp_clear */
945 0, /* tp_richcompare */
946 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000947 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000948 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000949 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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000958 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000959 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000960 PyObject_GC_Del, /* tp_free */
961};
962
963
964/* takewhile object **********************************************************/
965
966typedef struct {
967 PyObject_HEAD
968 PyObject *func;
969 PyObject *it;
970 long stop;
971} 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{
978 PyObject *func, *seq;
979 PyObject *it;
980 takewhileobject *lz;
981
Georg Brandlb84c1372007-01-21 10:28:43 +0000982 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000983 return NULL;
984
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000985 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
986 return NULL;
987
988 /* Get iterator. */
989 it = PyObject_GetIter(seq);
990 if (it == NULL)
991 return NULL;
992
993 /* 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;
1003
1004 return (PyObject *)lz;
1005}
1006
1007static void
1008takewhile_dealloc(takewhileobject *lz)
1009{
1010 PyObject_GC_UnTrack(lz);
1011 Py_XDECREF(lz->func);
1012 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001013 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001019 Py_VISIT(lz->it);
1020 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001021 return 0;
1022}
1023
1024static PyObject *
1025takewhile_next(takewhileobject *lz)
1026{
1027 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001028 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001029 long ok;
1030
1031 if (lz->stop == 1)
1032 return NULL;
1033
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001034 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001035 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036 if (item == NULL)
1037 return NULL;
1038
1039 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;
1051}
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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001060 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063 0, /* tp_itemsize */
1064 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001065 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066 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 */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001082 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001087 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001088 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001098 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001099 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100 PyObject_GC_Del, /* tp_free */
1101};
1102
1103
1104/* islice object ************************************************************/
1105
1106typedef struct {
1107 PyObject_HEAD
1108 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001109 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{
1120 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001121 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001123 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 isliceobject *lz;
1125
Georg Brandlb84c1372007-01-21 10:28:43 +00001126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001127 return NULL;
1128
Raymond Hettingerb2594052004-12-05 09:25:51 +00001129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 return NULL;
1131
Raymond Hettingerb2594052004-12-05 09:25:51 +00001132 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001135 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001140 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 return NULL;
1142 }
1143 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001144 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001146 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001147 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001148 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001149 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001150 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001155 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001156 return NULL;
1157 }
1158 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001160 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001162 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163 return NULL;
1164 }
1165
Raymond Hettingerb2594052004-12-05 09:25:51 +00001166 if (a3 != NULL) {
1167 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001168 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001174 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 return NULL;
1176 }
1177
1178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
1182
1183 /* 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;
1194
1195 return (PyObject *)lz;
1196}
1197
1198static void
1199islice_dealloc(isliceobject *lz)
1200{
1201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001203 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001209 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210 return 0;
1211}
1212
1213static PyObject *
1214islice_next(isliceobject *lz)
1215{
1216 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001217 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001218 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001219 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001221 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001222 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001224 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001230 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001232 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001233 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234 if (item == NULL)
1235 return NULL;
1236 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001237 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001238 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001239 if (lz->next < oldnext) /* Check for overflow */
1240 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241 return item;
1242}
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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001255 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001256 "itertools.islice", /* tp_name */
1257 sizeof(isliceobject), /* tp_basicsize */
1258 0, /* tp_itemsize */
1259 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001260 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261 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 */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001277 islice_doc, /* tp_doc */
1278 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279 0, /* tp_clear */
1280 0, /* tp_richcompare */
1281 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001282 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001283 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284 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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001293 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294 islice_new, /* tp_new */
1295 PyObject_GC_Del, /* tp_free */
1296};
1297
1298
1299/* starmap object ************************************************************/
1300
1301typedef struct {
1302 PyObject_HEAD
1303 PyObject *func;
1304 PyObject *it;
1305} 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{
1312 PyObject *func, *seq;
1313 PyObject *it;
1314 starmapobject *lz;
1315
Georg Brandlb84c1372007-01-21 10:28:43 +00001316 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001317 return NULL;
1318
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001319 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1320 return NULL;
1321
1322 /* Get iterator. */
1323 it = PyObject_GetIter(seq);
1324 if (it == NULL)
1325 return NULL;
1326
1327 /* 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;
1336
1337 return (PyObject *)lz;
1338}
1339
1340static void
1341starmap_dealloc(starmapobject *lz)
1342{
1343 PyObject_GC_UnTrack(lz);
1344 Py_XDECREF(lz->func);
1345 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00001346 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001352 Py_VISIT(lz->it);
1353 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354 return 0;
1355}
1356
1357static PyObject *
1358starmap_next(starmapobject *lz)
1359{
1360 PyObject *args;
1361 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001362 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001364 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00001365 args = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366 if (args == NULL)
1367 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001368 if (!PyTuple_CheckExact(args)) {
Raymond Hettinger47317092008-01-17 03:02:14 +00001369 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001370 Py_DECREF(args);
Raymond Hettinger47317092008-01-17 03:02:14 +00001371 if (newargs == NULL)
1372 return NULL;
1373 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001374 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375 result = PyObject_Call(lz->func, args, NULL);
1376 Py_DECREF(args);
1377 return result;
1378}
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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001387 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001388 "itertools.starmap", /* tp_name */
1389 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390 0, /* tp_itemsize */
1391 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001392 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001393 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 */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001409 starmap_doc, /* tp_doc */
1410 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001411 0, /* tp_clear */
1412 0, /* tp_richcompare */
1413 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001414 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001415 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001416 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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001425 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001426 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427 PyObject_GC_Del, /* tp_free */
1428};
1429
1430
1431/* imap object ************************************************************/
1432
1433typedef struct {
1434 PyObject_HEAD
1435 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436 PyObject *func;
1437} 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{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001444 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001446 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447
Georg Brandlb84c1372007-01-21 10:28:43 +00001448 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001449 return NULL;
1450
Raymond Hettinger96ef8112003-02-01 00:10:11 +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 }
1457
1458 iters = PyTuple_New(numargs-1);
1459 if (iters == NULL)
1460 return NULL;
1461
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462 for (i=1 ; i<numargs ; i++) {
1463 /* Get iterator. */
1464 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1465 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466 Py_DECREF(iters);
1467 return NULL;
1468 }
1469 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470 }
1471
1472 /* create imapobject structure */
1473 lz = (imapobject *)type->tp_alloc(type, 0);
1474 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475 Py_DECREF(iters);
1476 return NULL;
1477 }
1478 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479 func = PyTuple_GET_ITEM(args, 0);
1480 Py_INCREF(func);
1481 lz->func = func;
1482
1483 return (PyObject *)lz;
1484}
1485
1486static void
1487imap_dealloc(imapobject *lz)
1488{
1489 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490 Py_XDECREF(lz->iters);
1491 Py_XDECREF(lz->func);
Christian Heimese93237d2007-12-19 02:37:44 +00001492 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001498 Py_VISIT(lz->iters);
1499 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500 return 0;
1501}
1502
Raymond Hettinger2012f172003-02-07 05:32:58 +00001503/*
1504imap() 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
1514 2) In typical use cases for combining itertools, having one finite data
1515 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
1520 None.
1521
1522 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{
1531 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001532 PyObject *argtuple;
1533 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001534 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535
1536 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001537 argtuple = PyTuple_New(numargs);
1538 if (argtuple == NULL)
1539 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540
Raymond Hettingerf0c00242003-02-07 07:26:25 +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;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001546 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001547 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001548 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001549 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\
1560each of the iterables. Like map() except that it returns\n\
1561an 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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001566 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001567 "itertools.imap", /* tp_name */
1568 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001569 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 */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001593 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001594 (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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001604 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001605 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606 PyObject_GC_Del, /* tp_free */
1607};
1608
1609
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001610/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611
1612typedef struct {
1613 PyObject_HEAD
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001614 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
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001620static PyObject *
1621chain_new_internal(PyTypeObject *type, PyObject *source)
1622{
1623 chainobject *lz;
1624
1625 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;
1634}
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{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001639 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640
Georg Brandlb84c1372007-01-21 10:28:43 +00001641 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001642 return NULL;
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001643
1644 source = PyObject_GetIter(args);
1645 if (source == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001646 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001648 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{
1654 PyObject *source;
1655
1656 source = PyObject_GetIter(arg);
1657 if (source == NULL)
1658 return NULL;
1659
1660 return chain_new_internal(type, source);
1661}
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{
1666 PyObject_GC_UnTrack(lz);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001667 Py_XDECREF(lz->active);
1668 Py_XDECREF(lz->source);
Christian Heimese93237d2007-12-19 02:37:44 +00001669 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{
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001675 Py_VISIT(lz->source);
1676 Py_VISIT(lz->active);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001677 return 0;
1678}
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{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001683 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001685 if (lz->source == NULL)
1686 return NULL; /* already stopped */
1687
1688 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 */
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001693 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001694 lz->active = PyObject_GetIter(iterable);
Raymond Hettingerf1cca2b2008-03-04 22:29:44 +00001695 Py_DECREF(iterable);
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001696 if (lz->active == NULL) {
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001697 Py_CLEAR(lz->source);
1698 return NULL; /* input not iterable */
1699 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001700 }
Raymond Hettinger05bf6332008-02-28 22:30:42 +00001701 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[] = {
1728 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1729 chain_from_iterable_doc},
1730 {NULL, NULL} /* sentinel */
1731};
1732
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001733static PyTypeObject chain_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001734 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001735 "itertools.chain", /* tp_name */
1736 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737 0, /* tp_itemsize */
1738 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001739 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001740 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 */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001756 chain_doc, /* tp_doc */
1757 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758 0, /* tp_clear */
1759 0, /* tp_richcompare */
1760 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001761 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001762 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettingerb4cbc982008-02-28 22:46:41 +00001763 chain_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764 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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001772 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001773 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001774 PyObject_GC_Del, /* tp_free */
1775};
1776
1777
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001778/* product object ************************************************************/
1779
1780typedef struct {
1781 PyObject_HEAD
1782 PyObject *pools; /* tuple of pool tuples */
Raymond Hettinger532316d2008-02-23 04:03:50 +00001783 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{
1793 productobject *lz;
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001794 Py_ssize_t nargs, npools, repeat=1;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001795 PyObject *pools = NULL;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001796 Py_ssize_t *indices = NULL;
1797 Py_ssize_t i;
1798
Raymond Hettinger08ff6822008-02-29 02:21:48 +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
1816 assert(PyTuple_Check(args));
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001817 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1818 npools = nargs * repeat;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001819
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001820 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
Raymond Hettinger61024b92008-03-02 11:57:16 +00001821 if (indices == NULL) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001822 PyErr_NoMemory();
1823 goto error;
1824 }
1825
1826 pools = PyTuple_New(npools);
1827 if (pools == NULL)
1828 goto error;
1829
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001830 for (i=0; i < nargs ; ++i) {
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001831 PyObject *item = PyTuple_GET_ITEM(args, i);
1832 PyObject *pool = PySequence_Tuple(item);
1833 if (pool == NULL)
1834 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001835 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001836 indices[i] = 0;
1837 }
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001838 for ( ; i < npools; ++i) {
1839 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1840 Py_INCREF(pool);
1841 PyTuple_SET_ITEM(pools, i, pool);
Raymond Hettinger08ff6822008-02-29 02:21:48 +00001842 indices[i] = 0;
1843 }
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001844
1845 /* create productobject structure */
1846 lz = (productobject *)type->tp_alloc(type, 0);
Raymond Hettinger3bd77122008-02-27 01:08:04 +00001847 if (lz == NULL)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001848 goto error;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001849
1850 lz->pools = pools;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001851 lz->indices = indices;
1852 lz->result = NULL;
1853 lz->stopped = 0;
1854
1855 return (PyObject *)lz;
1856
1857error:
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001858 if (indices != NULL)
1859 PyMem_Free(indices);
1860 Py_XDECREF(pools);
1861 return NULL;
1862}
1863
1864static void
1865product_dealloc(productobject *lz)
1866{
1867 PyObject_GC_UnTrack(lz);
1868 Py_XDECREF(lz->pools);
1869 Py_XDECREF(lz->result);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001870 PyMem_Free(lz->indices);
1871 Py_TYPE(lz)->tp_free(lz);
1872}
1873
1874static int
1875product_traverse(productobject *lz, visitproc visit, void *arg)
1876{
1877 Py_VISIT(lz->pools);
1878 Py_VISIT(lz->result);
1879 return 0;
1880}
1881
1882static PyObject *
1883product_next(productobject *lz)
1884{
1885 PyObject *pool;
1886 PyObject *elem;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001887 PyObject *oldelem;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001888 PyObject *pools = lz->pools;
1889 PyObject *result = lz->result;
1890 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1891 Py_ssize_t i;
1892
1893 if (lz->stopped)
1894 return NULL;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001895
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001896 if (result == NULL) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001897 /* On the first pass, return an initial tuple filled with the
Raymond Hettingerd553d852008-03-04 04:17:08 +00001898 first element from each pool. */
Raymond Hettinger73d79632008-02-23 02:20:41 +00001899 result = PyTuple_New(npools);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001900 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);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001909 PyTuple_SET_ITEM(result, i, elem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001910 }
1911 } else {
1912 Py_ssize_t *indices = lz->indices;
Raymond Hettinger73d79632008-02-23 02:20:41 +00001913
1914 /* 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 */
Raymond Hettingere3fabd12008-03-02 12:02:19 +00001929 assert (npools==0 || Py_REFCNT(result) == 1);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001930
1931 /* Update the pool indices right-to-left. Only advance to the
1932 next pool when the previous one rolls-over */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001933 for (i=npools-1 ; i >= 0 ; i--) {
1934 pool = PyTuple_GET_ITEM(pools, i);
1935 indices[i]++;
Raymond Hettinger61024b92008-03-02 11:57:16 +00001936 if (indices[i] == PyTuple_GET_SIZE(pool)) {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001937 /* Roll-over and advance to next pool */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001938 indices[i] = 0;
1939 elem = PyTuple_GET_ITEM(pool, 0);
1940 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001941 oldelem = PyTuple_GET_ITEM(result, i);
1942 PyTuple_SET_ITEM(result, i, elem);
1943 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001944 } else {
Raymond Hettinger73d79632008-02-23 02:20:41 +00001945 /* No rollover. Just increment and stop here. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001946 elem = PyTuple_GET_ITEM(pool, indices[i]);
1947 Py_INCREF(elem);
Raymond Hettinger73d79632008-02-23 02:20:41 +00001948 oldelem = PyTuple_GET_ITEM(result, i);
1949 PyTuple_SET_ITEM(result, i, elem);
1950 Py_DECREF(oldelem);
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001951 break;
1952 }
1953 }
Raymond Hettinger73d79632008-02-23 02:20:41 +00001954
1955 /* If i is negative, then the indices have all rolled-over
1956 and we're done. */
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001957 if (i < 0)
Raymond Hettinger73d79632008-02-23 02:20:41 +00001958 goto empty;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001959 }
1960
Raymond Hettinger73d79632008-02-23 02:20:41 +00001961 Py_INCREF(result);
1962 return result;
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001963
1964empty:
1965 lz->stopped = 1;
1966 return NULL;
1967}
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 = {
1981 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 */
2022};
2023
2024
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002025/* combinations object ************************************************************/
2026
2027typedef struct {
2028 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 */
2034} combinationsobject;
2035
2036static PyTypeObject combinations_type;
2037
2038static PyObject *
2039combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2040{
2041 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};
2049
2050 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2051 &iterable, &r))
2052 return NULL;
2053
2054 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
2063 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2064 if (indices == NULL) {
2065 PyErr_NoMemory();
2066 goto error;
2067 }
2068
2069 for (i=0 ; i<r ; i++)
2070 indices[i] = i;
2071
2072 /* create combinationsobject structure */
2073 co = (combinationsobject *)type->tp_alloc(type, 0);
2074 if (co == NULL)
2075 goto error;
2076
2077 co->pool = pool;
2078 co->indices = indices;
2079 co->result = NULL;
2080 co->r = r;
Raymond Hettinger5b913e32009-01-08 06:39:04 +00002081 co->stopped = r > n ? 1 : 0;
Raymond Hettinger93e804d2008-02-26 23:40:50 +00002082
2083 return (PyObject *)co;
2084
2085error:
2086 if (indices != NULL)
2087 PyMem_Free(indices);
2088 Py_XDECREF(pool);
2089 return NULL;
2090}
2091
2092static void
2093combinations_dealloc(combinationsobject *co)
2094{
2095 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);
2100}
2101
2102static int
2103combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2104{
2105 Py_VISIT(co->pool);
2106 Py_VISIT(co->result);
2107 return 0;
2108}
2109
2110static PyObject *
2111combinations_next(combinationsobject *co)
2112{
2113 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;
2121
2122 if (co->stopped)
2123 return NULL;
2124
2125 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 }
Christian Heimescdddf182008-02-28 11:18:49 +00002152 /* 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
2158 /* 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 ;
2162
2163 /* 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;
2167
2168 /* 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;
2175
2176 /* 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 }
2187
2188 Py_INCREF(result);
2189 return result;
2190
2191empty:
2192 co->stopped = 1;
2193 return NULL;
2194}
2195
2196PyDoc_STRVAR(combinations_doc,
Georg Brandl8550d692008-06-10 12:46:39 +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 = {
2203 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 */
2244};
2245
2246
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002247/* permutations object ************************************************************
2248
2249def 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:
2258 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
2263 else:
2264 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
2270*/
2271
2272typedef struct {
2273 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 */
2280} permutationsobject;
2281
2282static PyTypeObject permutations_type;
2283
2284static PyObject *
2285permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2286{
2287 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};
2297
2298 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2299 &iterable, &robj))
2300 return NULL;
2301
2302 pool = PySequence_Tuple(iterable);
2303 if (pool == NULL)
2304 goto error;
2305 n = PyTuple_GET_SIZE(pool);
2306
2307 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
2318 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 }
2324
2325 for (i=0 ; i<n ; i++)
2326 indices[i] = i;
2327 for (i=0 ; i<r ; i++)
2328 cycles[i] = n - i;
2329
2330 /* create permutationsobject structure */
2331 po = (permutationsobject *)type->tp_alloc(type, 0);
2332 if (po == NULL)
2333 goto error;
2334
2335 po->pool = pool;
2336 po->indices = indices;
2337 po->cycles = cycles;
2338 po->result = NULL;
2339 po->r = r;
Raymond Hettinger5b913e32009-01-08 06:39:04 +00002340 po->stopped = r > n ? 1 : 0;
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002341
2342 return (PyObject *)po;
2343
2344error:
2345 if (indices != NULL)
2346 PyMem_Free(indices);
2347 if (cycles != NULL)
2348 PyMem_Free(cycles);
2349 Py_XDECREF(pool);
2350 return NULL;
2351}
2352
2353static void
2354permutations_dealloc(permutationsobject *po)
2355{
2356 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);
2362}
2363
2364static int
2365permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2366{
Raymond Hettinger6e3e4152008-03-05 21:04:32 +00002367 Py_VISIT(po->pool);
2368 Py_VISIT(po->result);
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00002369 return 0;
2370}
2371
2372static PyObject *
2373permutations_next(permutationsobject *po)
2374{
2375 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;
2384
2385 if (po->stopped)
2386 return NULL;
2387
2388 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;
2403
2404 /* Copy the previous result tuple or re-use it if available */
2405 if (Py_REFCNT(result) > 1) {
2406 PyObject *old_result = result;
2407 result = PyTuple_New(r);
2408 if (result == NULL)
2409 goto empty;
2410 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);
2420
2421 /* 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;
2436
2437 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;
2457
2458empty:
2459 po->stopped = 1;
2460 return NULL;
2461}
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 = {
2470 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 */
2511};
2512
2513
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002514/* ifilter object ************************************************************/
2515
2516typedef struct {
2517 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{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002527 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002528 PyObject *it;
2529 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002530
Georg Brandlb84c1372007-01-21 10:28:43 +00002531 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002532 return NULL;
2533
Raymond Hettinger60eca932003-02-09 06:40:58 +00002534 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002535 return NULL;
2536
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002537 /* Get iterator. */
2538 it = PyObject_GetIter(seq);
2539 if (it == NULL)
2540 return NULL;
2541
2542 /* 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
2552 return (PyObject *)lz;
2553}
2554
2555static void
2556ifilter_dealloc(ifilterobject *lz)
2557{
2558 PyObject_GC_UnTrack(lz);
2559 Py_XDECREF(lz->func);
2560 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002561 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002567 Py_VISIT(lz->it);
2568 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002569 return 0;
2570}
2571
2572static PyObject *
2573ifilter_next(ifilterobject *lz)
2574{
2575 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002576 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002577 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002578 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002579
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002580 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002581 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002582 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002583 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002584 if (item == NULL)
2585 return NULL;
2586
Facundo Batistab1d70e22008-02-25 23:46:02 +00002587 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002588 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 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002600 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002601 return item;
2602 Py_DECREF(item);
2603 }
2604}
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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002613 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002614 "itertools.ifilter", /* tp_name */
2615 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002616 0, /* tp_itemsize */
2617 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002618 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002619 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 */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002635 ifilter_doc, /* tp_doc */
2636 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002637 0, /* tp_clear */
2638 0, /* tp_richcompare */
2639 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002640 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002641 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002642 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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002651 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002652 ifilter_new, /* tp_new */
2653 PyObject_GC_Del, /* tp_free */
2654};
2655
2656
2657/* ifilterfalse object ************************************************************/
2658
2659typedef struct {
2660 PyObject_HEAD
2661 PyObject *func;
2662 PyObject *it;
2663} 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{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002670 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002671 PyObject *it;
2672 ifilterfalseobject *lz;
2673
Georg Brandlb84c1372007-01-21 10:28:43 +00002674 if (type == &ifilterfalse_type &&
2675 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002676 return NULL;
2677
Raymond Hettinger60eca932003-02-09 06:40:58 +00002678 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2679 return NULL;
2680
2681 /* Get iterator. */
2682 it = PyObject_GetIter(seq);
2683 if (it == NULL)
2684 return NULL;
2685
2686 /* 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;
2695
2696 return (PyObject *)lz;
2697}
2698
2699static void
2700ifilterfalse_dealloc(ifilterfalseobject *lz)
2701{
2702 PyObject_GC_UnTrack(lz);
2703 Py_XDECREF(lz->func);
2704 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002705 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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002711 Py_VISIT(lz->it);
2712 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002713 return 0;
2714}
2715
2716static PyObject *
2717ifilterfalse_next(ifilterfalseobject *lz)
2718{
2719 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002720 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002721 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002722 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002723
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002724 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002725 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002726 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002727 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002728 if (item == NULL)
2729 return NULL;
2730
Facundo Batistab1d70e22008-02-25 23:46:02 +00002731 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
Raymond Hettinger60eca932003-02-09 06:40:58 +00002732 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 }
2748}
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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002757 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002758 "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 */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002784 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002785 (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 */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002795 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002796 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002797 PyObject_GC_Del, /* tp_free */
2798};
2799
2800
2801/* count object ************************************************************/
2802
2803typedef struct {
2804 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002805 Py_ssize_t cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002806 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{
2814 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002815 Py_ssize_t cnt = 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002816 PyObject *cnt_arg = NULL;
2817 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002818
Georg Brandlb84c1372007-01-21 10:28:43 +00002819 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002820 return NULL;
2821
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002822 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002823 return NULL;
2824
Raymond Hettinger50e90e22007-10-04 00:20:27 +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 }
2838
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002839 /* create countobject structure */
2840 lz = (countobject *)PyObject_New(countobject, &count_type);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002841 if (lz == NULL) {
2842 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002843 return NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002844 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002845 lz->cnt = cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002846 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002847
2848 return (PyObject *)lz;
2849}
2850
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002851static void
2852count_dealloc(countobject *lz)
2853{
2854 Py_XDECREF(lz->long_cnt);
2855 PyObject_Del(lz);
2856}
2857
2858static PyObject *
2859count_nextlong(countobject *lz)
2860{
2861 static PyObject *one = NULL;
2862 PyObject *cnt;
2863 PyObject *stepped_up;
2864
2865 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;
2882}
2883
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002884static PyObject *
2885count_next(countobject *lz)
2886{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002887 if (lz->cnt == PY_SSIZE_T_MAX)
2888 return count_nextlong(lz);
Jack Diederich36234e82006-09-21 17:50:26 +00002889 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{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002895 PyObject *cnt_repr;
2896 PyObject *result;
2897
2898 if (lz->cnt != PY_SSIZE_T_MAX)
Gregory P. Smithdd96db62008-06-09 04:58:54 +00002899 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002900
2901 cnt_repr = PyObject_Repr(lz->long_cnt);
2902 if (cnt_repr == NULL)
2903 return NULL;
Gregory P. Smithdd96db62008-06-09 04:58:54 +00002904 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002905 Py_DECREF(cnt_repr);
2906 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002907}
2908
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002909PyDoc_STRVAR(count_doc,
2910"count([firstval]) --> count object\n\
2911\n\
2912Return a count object whose .next() method returns consecutive\n\
2913integers starting from zero or, if specified, from firstval.");
2914
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002915static PyTypeObject count_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002916 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002917 "itertools.count", /* tp_name */
2918 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002919 0, /* tp_itemsize */
2920 /* methods */
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002921 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002922 0, /* tp_print */
2923 0, /* tp_getattr */
2924 0, /* tp_setattr */
2925 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002926 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002927 0, /* tp_as_number */
2928 0, /* tp_as_sequence */
2929 0, /* tp_as_mapping */
2930 0, /* tp_hash */
2931 0, /* tp_call */
2932 0, /* tp_str */
2933 PyObject_GenericGetAttr, /* tp_getattro */
2934 0, /* tp_setattro */
2935 0, /* tp_as_buffer */
2936 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002937 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002938 0, /* tp_traverse */
2939 0, /* tp_clear */
2940 0, /* tp_richcompare */
2941 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002942 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002943 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002944 0, /* tp_methods */
2945 0, /* tp_members */
2946 0, /* tp_getset */
2947 0, /* tp_base */
2948 0, /* tp_dict */
2949 0, /* tp_descr_get */
2950 0, /* tp_descr_set */
2951 0, /* tp_dictoffset */
2952 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002953 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002954 count_new, /* tp_new */
2955};
2956
2957
2958/* izip object ************************************************************/
2959
2960#include "Python.h"
2961
2962typedef struct {
2963 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002964 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002965 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002966 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002967} izipobject;
2968
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002969static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002970
2971static PyObject *
2972izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2973{
2974 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002975 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002976 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002977 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002978 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002979
Georg Brandlb84c1372007-01-21 10:28:43 +00002980 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002981 return NULL;
2982
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002983 /* args must be a tuple */
2984 assert(PyTuple_Check(args));
2985
2986 /* obtain iterators */
2987 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002988 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002989 return NULL;
2990 for (i=0; i < tuplesize; ++i) {
2991 PyObject *item = PyTuple_GET_ITEM(args, i);
2992 PyObject *it = PyObject_GetIter(item);
2993 if (it == NULL) {
2994 if (PyErr_ExceptionMatches(PyExc_TypeError))
2995 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002996 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002997 i+1);
2998 Py_DECREF(ittuple);
2999 return NULL;
3000 }
3001 PyTuple_SET_ITEM(ittuple, i, it);
3002 }
3003
Raymond Hettinger2012f172003-02-07 05:32:58 +00003004 /* create a result holder */
3005 result = PyTuple_New(tuplesize);
3006 if (result == NULL) {
3007 Py_DECREF(ittuple);
3008 return NULL;
3009 }
3010 for (i=0 ; i < tuplesize ; i++) {
3011 Py_INCREF(Py_None);
3012 PyTuple_SET_ITEM(result, i, Py_None);
3013 }
3014
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003015 /* create izipobject structure */
3016 lz = (izipobject *)type->tp_alloc(type, 0);
3017 if (lz == NULL) {
3018 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003019 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003020 return NULL;
3021 }
3022 lz->ittuple = ittuple;
3023 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003024 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003025
3026 return (PyObject *)lz;
3027}
3028
3029static void
3030izip_dealloc(izipobject *lz)
3031{
3032 PyObject_GC_UnTrack(lz);
3033 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003034 Py_XDECREF(lz->result);
Christian Heimese93237d2007-12-19 02:37:44 +00003035 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003036}
3037
3038static int
3039izip_traverse(izipobject *lz, visitproc visit, void *arg)
3040{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003041 Py_VISIT(lz->ittuple);
3042 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003043 return 0;
3044}
3045
3046static PyObject *
3047izip_next(izipobject *lz)
3048{
Jack Diederich6c433a92006-05-26 11:15:17 +00003049 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00003050 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00003051 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003052 PyObject *it;
3053 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003054 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003055
Raymond Hettingerb5a42082003-08-08 05:10:41 +00003056 if (tuplesize == 0)
3057 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003058 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003059 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003060 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003061 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00003062 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003063 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003064 if (item == NULL) {
3065 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003066 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00003067 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003068 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003069 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00003070 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00003071 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00003072 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00003073 result = PyTuple_New(tuplesize);
3074 if (result == NULL)
3075 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003076 for (i=0 ; i < tuplesize ; i++) {
3077 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00003078 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003079 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00003080 if (item == NULL) {
3081 Py_DECREF(result);
3082 return NULL;
3083 }
3084 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003085 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003086 }
3087 return result;
3088}
3089
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003090PyDoc_STRVAR(izip_doc,
3091"izip(iter1 [,iter2 [...]]) --> izip object\n\
3092\n\
3093Return a izip object whose .next() method returns a tuple where\n\
3094the i-th element comes from the i-th iterable argument. The .next()\n\
3095method continues until the shortest iterable in the argument sequence\n\
3096is exhausted and then it raises StopIteration. Works like the zip()\n\
3097function but consumes less memory by returning an iterator instead of\n\
3098a list.");
3099
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003100static PyTypeObject izip_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003101 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003102 "itertools.izip", /* tp_name */
3103 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003104 0, /* tp_itemsize */
3105 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003106 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003107 0, /* tp_print */
3108 0, /* tp_getattr */
3109 0, /* tp_setattr */
3110 0, /* tp_compare */
3111 0, /* tp_repr */
3112 0, /* tp_as_number */
3113 0, /* tp_as_sequence */
3114 0, /* tp_as_mapping */
3115 0, /* tp_hash */
3116 0, /* tp_call */
3117 0, /* tp_str */
3118 PyObject_GenericGetAttr, /* tp_getattro */
3119 0, /* tp_setattro */
3120 0, /* tp_as_buffer */
3121 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3122 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003123 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003124 (traverseproc)izip_traverse, /* tp_traverse */
3125 0, /* tp_clear */
3126 0, /* tp_richcompare */
3127 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003128 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003129 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003130 0, /* tp_methods */
3131 0, /* tp_members */
3132 0, /* tp_getset */
3133 0, /* tp_base */
3134 0, /* tp_dict */
3135 0, /* tp_descr_get */
3136 0, /* tp_descr_set */
3137 0, /* tp_dictoffset */
3138 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003139 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003140 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003141 PyObject_GC_Del, /* tp_free */
3142};
3143
3144
3145/* repeat object ************************************************************/
3146
3147typedef struct {
3148 PyObject_HEAD
3149 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003150 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003151} repeatobject;
3152
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003153static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003154
3155static PyObject *
3156repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3157{
3158 repeatobject *ro;
3159 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00003160 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003161
Georg Brandlb84c1372007-01-21 10:28:43 +00003162 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00003163 return NULL;
3164
Jack Diederich6c433a92006-05-26 11:15:17 +00003165 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003166 return NULL;
3167
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00003168 if (PyTuple_Size(args) == 2 && cnt < 0)
3169 cnt = 0;
3170
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003171 ro = (repeatobject *)type->tp_alloc(type, 0);
3172 if (ro == NULL)
3173 return NULL;
3174 Py_INCREF(element);
3175 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003176 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003177 return (PyObject *)ro;
3178}
3179
3180static void
3181repeat_dealloc(repeatobject *ro)
3182{
3183 PyObject_GC_UnTrack(ro);
3184 Py_XDECREF(ro->element);
Christian Heimese93237d2007-12-19 02:37:44 +00003185 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003186}
3187
3188static int
3189repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3190{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00003191 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003192 return 0;
3193}
3194
3195static PyObject *
3196repeat_next(repeatobject *ro)
3197{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003198 if (ro->cnt == 0)
3199 return NULL;
3200 if (ro->cnt > 0)
3201 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003202 Py_INCREF(ro->element);
3203 return ro->element;
3204}
3205
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003206static PyObject *
3207repeat_repr(repeatobject *ro)
3208{
3209 PyObject *result, *objrepr;
3210
3211 objrepr = PyObject_Repr(ro->element);
3212 if (objrepr == NULL)
3213 return NULL;
3214
3215 if (ro->cnt == -1)
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003216 result = PyString_FromFormat("repeat(%s)",
3217 PyString_AS_STRING(objrepr));
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003218 else
Gregory P. Smithdd96db62008-06-09 04:58:54 +00003219 result = PyString_FromFormat("repeat(%s, %zd)",
3220 PyString_AS_STRING(objrepr), ro->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003221 Py_DECREF(objrepr);
3222 return result;
3223}
3224
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003225static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003226repeat_len(repeatobject *ro)
3227{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003228 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003229 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003230 return NULL;
3231 }
Jack Diederich6c433a92006-05-26 11:15:17 +00003232 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003233}
3234
Armin Rigof5b3e362006-02-11 21:32:43 +00003235PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003236
3237static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00003238 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003239 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00003240};
3241
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003242PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003243"repeat(element [,times]) -> create an iterator which returns the element\n\
3244for the specified number of times. If not specified, returns the element\n\
3245endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003246
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00003247static PyTypeObject repeat_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003248 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003249 "itertools.repeat", /* tp_name */
3250 sizeof(repeatobject), /* tp_basicsize */
3251 0, /* tp_itemsize */
3252 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003253 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003254 0, /* tp_print */
3255 0, /* tp_getattr */
3256 0, /* tp_setattr */
3257 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00003258 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003259 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003260 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003261 0, /* tp_as_mapping */
3262 0, /* tp_hash */
3263 0, /* tp_call */
3264 0, /* tp_str */
3265 PyObject_GenericGetAttr, /* tp_getattro */
3266 0, /* tp_setattro */
3267 0, /* tp_as_buffer */
3268 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3269 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003270 repeat_doc, /* tp_doc */
3271 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003272 0, /* tp_clear */
3273 0, /* tp_richcompare */
3274 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00003275 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00003276 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00003277 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003278 0, /* tp_members */
3279 0, /* tp_getset */
3280 0, /* tp_base */
3281 0, /* tp_dict */
3282 0, /* tp_descr_get */
3283 0, /* tp_descr_set */
3284 0, /* tp_dictoffset */
3285 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00003286 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003287 repeat_new, /* tp_new */
3288 PyObject_GC_Del, /* tp_free */
3289};
3290
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003291/* iziplongest object ************************************************************/
3292
3293#include "Python.h"
3294
3295typedef struct {
3296 PyObject_HEAD
3297 Py_ssize_t tuplesize;
3298 Py_ssize_t numactive;
3299 PyObject *ittuple; /* tuple of iterators */
3300 PyObject *result;
3301 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003302} iziplongestobject;
3303
3304static PyTypeObject iziplongest_type;
3305
3306static PyObject *
3307izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3308{
3309 iziplongestobject *lz;
3310 Py_ssize_t i;
3311 PyObject *ittuple; /* tuple of iterators */
3312 PyObject *result;
3313 PyObject *fillvalue = Py_None;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003314 Py_ssize_t tuplesize = PySequence_Length(args);
3315
3316 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3317 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3318 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3319 PyErr_SetString(PyExc_TypeError,
3320 "izip_longest() got an unexpected keyword argument");
3321 return NULL;
3322 }
3323 }
3324
3325 /* args must be a tuple */
3326 assert(PyTuple_Check(args));
3327
3328 /* obtain iterators */
3329 ittuple = PyTuple_New(tuplesize);
3330 if (ittuple == NULL)
3331 return NULL;
3332 for (i=0; i < tuplesize; ++i) {
3333 PyObject *item = PyTuple_GET_ITEM(args, i);
3334 PyObject *it = PyObject_GetIter(item);
3335 if (it == NULL) {
3336 if (PyErr_ExceptionMatches(PyExc_TypeError))
3337 PyErr_Format(PyExc_TypeError,
3338 "izip_longest argument #%zd must support iteration",
3339 i+1);
3340 Py_DECREF(ittuple);
3341 return NULL;
3342 }
3343 PyTuple_SET_ITEM(ittuple, i, it);
3344 }
3345
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003346 /* create a result holder */
3347 result = PyTuple_New(tuplesize);
3348 if (result == NULL) {
3349 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003350 return NULL;
3351 }
3352 for (i=0 ; i < tuplesize ; i++) {
3353 Py_INCREF(Py_None);
3354 PyTuple_SET_ITEM(result, i, Py_None);
3355 }
3356
3357 /* create iziplongestobject structure */
3358 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3359 if (lz == NULL) {
3360 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003361 Py_DECREF(result);
3362 return NULL;
3363 }
3364 lz->ittuple = ittuple;
3365 lz->tuplesize = tuplesize;
3366 lz->numactive = tuplesize;
3367 lz->result = result;
3368 Py_INCREF(fillvalue);
3369 lz->fillvalue = fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003370 return (PyObject *)lz;
3371}
3372
3373static void
3374izip_longest_dealloc(iziplongestobject *lz)
3375{
3376 PyObject_GC_UnTrack(lz);
3377 Py_XDECREF(lz->ittuple);
3378 Py_XDECREF(lz->result);
3379 Py_XDECREF(lz->fillvalue);
Christian Heimese93237d2007-12-19 02:37:44 +00003380 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003381}
3382
3383static int
3384izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3385{
3386 Py_VISIT(lz->ittuple);
3387 Py_VISIT(lz->result);
3388 Py_VISIT(lz->fillvalue);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003389 return 0;
3390}
3391
3392static PyObject *
3393izip_longest_next(iziplongestobject *lz)
3394{
3395 Py_ssize_t i;
3396 Py_ssize_t tuplesize = lz->tuplesize;
3397 PyObject *result = lz->result;
3398 PyObject *it;
3399 PyObject *item;
3400 PyObject *olditem;
3401
3402 if (tuplesize == 0)
3403 return NULL;
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003404 if (lz->numactive == 0)
3405 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00003406 if (Py_REFCNT(result) == 1) {
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003407 Py_INCREF(result);
3408 for (i=0 ; i < tuplesize ; i++) {
3409 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003410 if (it == NULL) {
3411 Py_INCREF(lz->fillvalue);
3412 item = lz->fillvalue;
3413 } else {
3414 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003415 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003416 if (item == NULL) {
3417 lz->numactive -= 1;
3418 if (lz->numactive == 0) {
3419 Py_DECREF(result);
3420 return NULL;
3421 } else {
3422 Py_INCREF(lz->fillvalue);
3423 item = lz->fillvalue;
3424 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3425 Py_DECREF(it);
3426 }
3427 }
3428 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003429 olditem = PyTuple_GET_ITEM(result, i);
3430 PyTuple_SET_ITEM(result, i, item);
3431 Py_DECREF(olditem);
3432 }
3433 } else {
3434 result = PyTuple_New(tuplesize);
3435 if (result == NULL)
3436 return NULL;
3437 for (i=0 ; i < tuplesize ; i++) {
3438 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003439 if (it == NULL) {
3440 Py_INCREF(lz->fillvalue);
3441 item = lz->fillvalue;
3442 } else {
3443 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00003444 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00003445 if (item == NULL) {
3446 lz->numactive -= 1;
3447 if (lz->numactive == 0) {
3448 Py_DECREF(result);
3449 return NULL;
3450 } else {
3451 Py_INCREF(lz->fillvalue);
3452 item = lz->fillvalue;
3453 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3454 Py_DECREF(it);
3455 }
3456 }
3457 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003458 PyTuple_SET_ITEM(result, i, item);
3459 }
3460 }
3461 return result;
3462}
3463
3464PyDoc_STRVAR(izip_longest_doc,
3465"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3466\n\
3467Return an izip_longest object whose .next() method returns a tuple where\n\
3468the i-th element comes from the i-th iterable argument. The .next()\n\
3469method continues until the longest iterable in the argument sequence\n\
3470is exhausted and then it raises StopIteration. When the shorter iterables\n\
3471are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3472defaults to None or can be specified by a keyword argument.\n\
3473");
3474
3475static PyTypeObject iziplongest_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00003476 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003477 "itertools.izip_longest", /* tp_name */
3478 sizeof(iziplongestobject), /* tp_basicsize */
3479 0, /* tp_itemsize */
3480 /* methods */
3481 (destructor)izip_longest_dealloc, /* tp_dealloc */
3482 0, /* tp_print */
3483 0, /* tp_getattr */
3484 0, /* tp_setattr */
3485 0, /* tp_compare */
3486 0, /* tp_repr */
3487 0, /* tp_as_number */
3488 0, /* tp_as_sequence */
3489 0, /* tp_as_mapping */
3490 0, /* tp_hash */
3491 0, /* tp_call */
3492 0, /* tp_str */
3493 PyObject_GenericGetAttr, /* tp_getattro */
3494 0, /* tp_setattro */
3495 0, /* tp_as_buffer */
3496 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3497 Py_TPFLAGS_BASETYPE, /* tp_flags */
3498 izip_longest_doc, /* tp_doc */
3499 (traverseproc)izip_longest_traverse, /* tp_traverse */
3500 0, /* tp_clear */
3501 0, /* tp_richcompare */
3502 0, /* tp_weaklistoffset */
3503 PyObject_SelfIter, /* tp_iter */
3504 (iternextfunc)izip_longest_next, /* tp_iternext */
3505 0, /* tp_methods */
3506 0, /* tp_members */
3507 0, /* tp_getset */
3508 0, /* tp_base */
3509 0, /* tp_dict */
3510 0, /* tp_descr_get */
3511 0, /* tp_descr_set */
3512 0, /* tp_dictoffset */
3513 0, /* tp_init */
3514 0, /* tp_alloc */
3515 izip_longest_new, /* tp_new */
3516 PyObject_GC_Del, /* tp_free */
3517};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003518
3519/* module level code ********************************************************/
3520
3521PyDoc_STRVAR(module_doc,
3522"Functional tools for creating and using iterators.\n\
3523\n\
3524Infinite iterators:\n\
3525count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003526cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00003527repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003528\n\
3529Iterators terminating on the shortest input sequence:\n\
3530izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettingerd36862c2007-02-21 05:20:38 +00003531izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00003532ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3533ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003534islice(seq, [start,] stop [, step]) --> elements from\n\
3535 seq[start:stop:step]\n\
3536imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3537starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00003538tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003539chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003540takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3541dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003542groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003543");
3544
3545
Raymond Hettingerad983e72003-11-12 14:32:26 +00003546static PyMethodDef module_methods[] = {
3547 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3548 {NULL, NULL} /* sentinel */
3549};
3550
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003551PyMODINIT_FUNC
3552inititertools(void)
3553{
Raymond Hettinger60eca932003-02-09 06:40:58 +00003554 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003555 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003556 char *name;
3557 PyTypeObject *typelist[] = {
Raymond Hettinger93e804d2008-02-26 23:40:50 +00003558 &combinations_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003559 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003560 &dropwhile_type,
3561 &takewhile_type,
3562 &islice_type,
3563 &starmap_type,
3564 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003565 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003566 &ifilter_type,
3567 &ifilterfalse_type,
3568 &count_type,
3569 &izip_type,
Raymond Hettinger50986cc2008-02-22 03:16:42 +00003570 &iziplongest_type,
Raymond Hettinger66f91ea2008-03-05 20:59:58 +00003571 &permutations_type,
David Wolever2724ab92008-03-19 02:35:45 +00003572 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003573 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003574 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003575 NULL
3576 };
3577
Christian Heimese93237d2007-12-19 02:37:44 +00003578 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003579 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003580 if (m == NULL)
3581 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003582
Raymond Hettinger60eca932003-02-09 06:40:58 +00003583 for (i=0 ; typelist[i] != NULL ; i++) {
3584 if (PyType_Ready(typelist[i]) < 0)
3585 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003586 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003587 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003588 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003589 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003590 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003591
3592 if (PyType_Ready(&teedataobject_type) < 0)
3593 return;
3594 if (PyType_Ready(&tee_type) < 0)
3595 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003596 if (PyType_Ready(&_grouper_type) < 0)
3597 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003598}