blob: 34b37be43ff6369256b298418a21ed038ec78150 [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{
29 static char *kwargs[] = {"iterable", "key", NULL};
30 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);
62 gbo->ob_type->tp_free(gbo);
63}
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 = {
142 PyObject_HEAD_INIT(NULL)
143 0, /* ob_size */
144 "itertools.groupby", /* tp_name */
145 sizeof(groupbyobject), /* tp_basicsize */
146 0, /* tp_itemsize */
147 /* methods */
148 (destructor)groupby_dealloc, /* tp_dealloc */
149 0, /* tp_print */
150 0, /* tp_getattr */
151 0, /* tp_setattr */
152 0, /* tp_compare */
153 0, /* tp_repr */
154 0, /* tp_as_number */
155 0, /* tp_as_sequence */
156 0, /* tp_as_mapping */
157 0, /* tp_hash */
158 0, /* tp_call */
159 0, /* tp_str */
160 PyObject_GenericGetAttr, /* tp_getattro */
161 0, /* tp_setattro */
162 0, /* tp_as_buffer */
163 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
164 Py_TPFLAGS_BASETYPE, /* tp_flags */
165 groupby_doc, /* tp_doc */
166 (traverseproc)groupby_traverse, /* tp_traverse */
167 0, /* tp_clear */
168 0, /* tp_richcompare */
169 0, /* tp_weaklistoffset */
170 PyObject_SelfIter, /* tp_iter */
171 (iternextfunc)groupby_next, /* tp_iternext */
172 0, /* tp_methods */
173 0, /* tp_members */
174 0, /* tp_getset */
175 0, /* tp_base */
176 0, /* tp_dict */
177 0, /* tp_descr_get */
178 0, /* tp_descr_set */
179 0, /* tp_dictoffset */
180 0, /* tp_init */
181 0, /* tp_alloc */
182 groupby_new, /* tp_new */
183 PyObject_GC_Del, /* tp_free */
184};
185
186
187/* _grouper object (internal) ************************************************/
188
189typedef struct {
190 PyObject_HEAD
191 PyObject *parent;
192 PyObject *tgtkey;
193} _grouperobject;
194
195static PyTypeObject _grouper_type;
196
197static PyObject *
198_grouper_create(groupbyobject *parent, PyObject *tgtkey)
199{
200 _grouperobject *igo;
201
202 igo = PyObject_New(_grouperobject, &_grouper_type);
203 if (igo == NULL)
204 return NULL;
205 igo->parent = (PyObject *)parent;
206 Py_INCREF(parent);
207 igo->tgtkey = tgtkey;
208 Py_INCREF(tgtkey);
209
210 return (PyObject *)igo;
211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
216 Py_DECREF(igo->parent);
217 Py_DECREF(igo->tgtkey);
218 PyObject_Del(igo);
219}
220
221static PyObject *
222_grouper_next(_grouperobject *igo)
223{
224 groupbyobject *gbo = (groupbyobject *)igo->parent;
225 PyObject *newvalue, *newkey, *r;
226 int rcmp;
227
228 if (gbo->currvalue == NULL) {
229 newvalue = PyIter_Next(gbo->it);
230 if (newvalue == NULL)
231 return NULL;
232
233 if (gbo->keyfunc == Py_None) {
234 newkey = newvalue;
235 Py_INCREF(newvalue);
236 } else {
237 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
238 newvalue, NULL);
239 if (newkey == NULL) {
240 Py_DECREF(newvalue);
241 return NULL;
242 }
243 }
244
245 assert(gbo->currkey == NULL);
246 gbo->currkey = newkey;
247 gbo->currvalue = newvalue;
248 }
249
250 assert(gbo->currkey != NULL);
251 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
252 if (rcmp <= 0)
253 /* got any error or current group is end */
254 return NULL;
255
256 r = gbo->currvalue;
257 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000258 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000259
260 return r;
261}
262
263static PyTypeObject _grouper_type = {
264 PyObject_HEAD_INIT(NULL)
265 0, /* ob_size */
266 "itertools._grouper", /* tp_name */
267 sizeof(_grouperobject), /* tp_basicsize */
268 0, /* tp_itemsize */
269 /* methods */
270 (destructor)_grouper_dealloc, /* tp_dealloc */
271 0, /* tp_print */
272 0, /* tp_getattr */
273 0, /* tp_setattr */
274 0, /* tp_compare */
275 0, /* tp_repr */
276 0, /* tp_as_number */
277 0, /* tp_as_sequence */
278 0, /* tp_as_mapping */
279 0, /* tp_hash */
280 0, /* tp_call */
281 0, /* tp_str */
282 PyObject_GenericGetAttr, /* tp_getattro */
283 0, /* tp_setattro */
284 0, /* tp_as_buffer */
285 Py_TPFLAGS_DEFAULT, /* tp_flags */
286 0, /* tp_doc */
287 0, /* tp_traverse */
288 0, /* tp_clear */
289 0, /* tp_richcompare */
290 0, /* tp_weaklistoffset */
291 PyObject_SelfIter, /* tp_iter */
292 (iternextfunc)_grouper_next, /* tp_iternext */
293 0, /* tp_methods */
294 0, /* tp_members */
295 0, /* tp_getset */
296 0, /* tp_base */
297 0, /* tp_dict */
298 0, /* tp_descr_get */
299 0, /* tp_descr_set */
300 0, /* tp_dictoffset */
301 0, /* tp_init */
302 0, /* tp_alloc */
303 0, /* tp_new */
304 PyObject_Del, /* tp_free */
305};
306
307
308
Raymond Hettingerad983e72003-11-12 14:32:26 +0000309/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000310
Raymond Hettingerad983e72003-11-12 14:32:26 +0000311/* The teedataobject pre-allocates space for LINKCELLS number of objects.
312 To help the object fit neatly inside cache lines (space for 16 to 32
313 pointers), the value should be a multiple of 16 minus space for
314 the other structure members including PyHEAD overhead. The larger the
315 value, the less memory overhead per object and the less time spent
316 allocating/deallocating new links. The smaller the number, the less
317 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000320
321typedef struct {
322 PyObject_HEAD
323 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000324 int numread;
325 PyObject *nextlink;
326 PyObject *(values[LINKCELLS]);
327} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
330 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000331 teedataobject *dataobj;
332 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000333 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335
Raymond Hettingerad983e72003-11-12 14:32:26 +0000336static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000337
338static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000340{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000341 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000342
Raymond Hettingerad983e72003-11-12 14:32:26 +0000343 tdo = PyObject_New(teedataobject, &teedataobject_type);
344 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000345 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000346
347 tdo->numread = 0;
348 tdo->nextlink = NULL;
349 Py_INCREF(it);
350 tdo->it = it;
351 return (PyObject *)tdo;
352}
353
354static PyObject *
355teedataobject_jumplink(teedataobject *tdo)
356{
357 if (tdo->nextlink == NULL)
358 tdo->nextlink = teedataobject_new(tdo->it);
359 Py_INCREF(tdo->nextlink);
360 return tdo->nextlink;
361}
362
363static PyObject *
364teedataobject_getitem(teedataobject *tdo, int i)
365{
366 PyObject *value;
367
368 assert(i < LINKCELLS);
369 if (i < tdo->numread)
370 value = tdo->values[i];
371 else {
372 /* this is the lead iterator, so fetch more data */
373 assert(i == tdo->numread);
374 value = PyIter_Next(tdo->it);
375 if (value == NULL)
376 return NULL;
377 tdo->numread++;
378 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000379 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000380 Py_INCREF(value);
381 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000382}
383
384static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000385teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000386{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000387 int i;
Raymond Hettinger45143692003-10-25 06:37:47 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 for (i=0 ; i<tdo->numread ; i++)
390 Py_DECREF(tdo->values[i]);
391 Py_XDECREF(tdo->it);
392 Py_XDECREF(tdo->nextlink);
393 PyObject_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000394}
395
Raymond Hettingerad983e72003-11-12 14:32:26 +0000396PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000397
Raymond Hettingerad983e72003-11-12 14:32:26 +0000398static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000399 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000401 "itertools.tee_dataobject", /* tp_name */
402 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000403 0, /* tp_itemsize */
404 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000405 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000406 0, /* tp_print */
407 0, /* tp_getattr */
408 0, /* tp_setattr */
409 0, /* tp_compare */
410 0, /* tp_repr */
411 0, /* tp_as_number */
412 0, /* tp_as_sequence */
413 0, /* tp_as_mapping */
414 0, /* tp_hash */
415 0, /* tp_call */
416 0, /* tp_str */
417 PyObject_GenericGetAttr, /* tp_getattro */
418 0, /* tp_setattro */
419 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000420 Py_TPFLAGS_DEFAULT, /* tp_flags */
421 teedataobject_doc, /* tp_doc */
422 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000423};
424
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000425
426static PyTypeObject tee_type;
427
428static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000429tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000430{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431 PyObject *value, *link;
432
433 if (to->index >= LINKCELLS) {
434 link = teedataobject_jumplink(to->dataobj);
435 Py_XDECREF(to->dataobj);
436 to->dataobj = (teedataobject *)link;
437 to->index = 0;
438 }
439 value = teedataobject_getitem(to->dataobj, to->index);
440 if (value == NULL)
441 return NULL;
442 to->index++;
443 return value;
444}
445
446static PyObject *
447tee_copy(teeobject *to)
448{
449 teeobject *newto;
450
451 newto = PyObject_New(teeobject, &tee_type);
452 if (newto == NULL)
453 return NULL;
454 Py_INCREF(to->dataobj);
455 newto->dataobj = to->dataobj;
456 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000457 newto->weakreflist = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000458 return (PyObject *)newto;
459}
460
461PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
462
463static PyObject *
464tee_fromiterable(PyObject *iterable)
465{
466 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000467 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000468
469 it = PyObject_GetIter(iterable);
470 if (it == NULL)
471 return NULL;
472 if (PyObject_TypeCheck(it, &tee_type)) {
473 to = (teeobject *)tee_copy((teeobject *)it);
474 goto done;
475 }
476
477 to = PyObject_New(teeobject, &tee_type);
478 if (to == NULL)
479 goto done;
480 to->dataobj = (teedataobject *)teedataobject_new(it);
481 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000482 to->weakreflist = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000483done:
484 Py_XDECREF(it);
485 return (PyObject *)to;
486}
487
488static PyObject *
489tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
490{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000491 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000492
493 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
494 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000495 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000496}
497
498static void
499tee_dealloc(teeobject *to)
500{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000501 if (to->weakreflist != NULL)
502 PyObject_ClearWeakRefs((PyObject *) to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000503 Py_XDECREF(to->dataobj);
504 PyObject_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000505}
506
Raymond Hettingerad983e72003-11-12 14:32:26 +0000507PyDoc_STRVAR(teeobject_doc,
508"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000509
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510static PyMethodDef tee_methods[] = {
511 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
512 {NULL, NULL} /* sentinel */
513};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000514
515static PyTypeObject tee_type = {
516 PyObject_HEAD_INIT(NULL)
517 0, /* ob_size */
518 "itertools.tee", /* tp_name */
519 sizeof(teeobject), /* tp_basicsize */
520 0, /* tp_itemsize */
521 /* methods */
522 (destructor)tee_dealloc, /* tp_dealloc */
523 0, /* tp_print */
524 0, /* tp_getattr */
525 0, /* tp_setattr */
526 0, /* tp_compare */
527 0, /* tp_repr */
528 0, /* tp_as_number */
529 0, /* tp_as_sequence */
530 0, /* tp_as_mapping */
531 0, /* tp_hash */
532 0, /* tp_call */
533 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000534 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000535 0, /* tp_setattro */
536 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000537 Py_TPFLAGS_DEFAULT, /* tp_flags */
538 teeobject_doc, /* tp_doc */
539 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000540 0, /* tp_clear */
541 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000542 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000543 PyObject_SelfIter, /* tp_iter */
544 (iternextfunc)tee_next, /* tp_iternext */
545 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000546 0, /* tp_members */
547 0, /* tp_getset */
548 0, /* tp_base */
549 0, /* tp_dict */
550 0, /* tp_descr_get */
551 0, /* tp_descr_set */
552 0, /* tp_dictoffset */
553 0, /* tp_init */
554 0, /* tp_alloc */
555 tee_new, /* tp_new */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000556 PyObject_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000557};
558
Raymond Hettingerad983e72003-11-12 14:32:26 +0000559static PyObject *
560tee(PyObject *self, PyObject *args)
561{
562 int i, n=2;
563 PyObject *it, *iterable, *copyable, *result;
564
565 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
566 return NULL;
567 result = PyTuple_New(n);
568 if (result == NULL)
569 return NULL;
570 if (n == 0)
571 return result;
572 it = PyObject_GetIter(iterable);
573 if (it == NULL) {
574 Py_DECREF(result);
575 return NULL;
576 }
577 if (!PyObject_HasAttrString(it, "__copy__")) {
578 copyable = tee_fromiterable(it);
579 Py_DECREF(it);
580 if (copyable == NULL) {
581 Py_DECREF(result);
582 return NULL;
583 }
584 } else
585 copyable = it;
586 PyTuple_SET_ITEM(result, 0, copyable);
587 for (i=1 ; i<n ; i++) {
588 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
589 if (copyable == NULL) {
590 Py_DECREF(result);
591 return NULL;
592 }
593 PyTuple_SET_ITEM(result, i, copyable);
594 }
595 return result;
596}
597
598PyDoc_STRVAR(tee_doc,
599"tee(iterable, n=2) --> tuple of n independent iterators.");
600
601
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000602/* cycle object **********************************************************/
603
604typedef struct {
605 PyObject_HEAD
606 PyObject *it;
607 PyObject *saved;
608 int firstpass;
609} cycleobject;
610
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000611static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000612
613static PyObject *
614cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
615{
616 PyObject *it;
617 PyObject *iterable;
618 PyObject *saved;
619 cycleobject *lz;
620
Georg Brandl02c42872005-08-26 06:42:30 +0000621 if (!_PyArg_NoKeywords("cycle()", kwds))
622 return NULL;
623
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000624 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
625 return NULL;
626
627 /* Get iterator. */
628 it = PyObject_GetIter(iterable);
629 if (it == NULL)
630 return NULL;
631
632 saved = PyList_New(0);
633 if (saved == NULL) {
634 Py_DECREF(it);
635 return NULL;
636 }
637
638 /* create cycleobject structure */
639 lz = (cycleobject *)type->tp_alloc(type, 0);
640 if (lz == NULL) {
641 Py_DECREF(it);
642 Py_DECREF(saved);
643 return NULL;
644 }
645 lz->it = it;
646 lz->saved = saved;
647 lz->firstpass = 0;
648
649 return (PyObject *)lz;
650}
651
652static void
653cycle_dealloc(cycleobject *lz)
654{
655 PyObject_GC_UnTrack(lz);
656 Py_XDECREF(lz->saved);
657 Py_XDECREF(lz->it);
658 lz->ob_type->tp_free(lz);
659}
660
661static int
662cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
663{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000664 Py_VISIT(lz->it);
665 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000666 return 0;
667}
668
669static PyObject *
670cycle_next(cycleobject *lz)
671{
672 PyObject *item;
673 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000674 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000675
676 while (1) {
677 item = PyIter_Next(lz->it);
678 if (item != NULL) {
679 if (!lz->firstpass)
680 PyList_Append(lz->saved, item);
681 return item;
682 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000683 if (PyErr_Occurred()) {
684 if (PyErr_ExceptionMatches(PyExc_StopIteration))
685 PyErr_Clear();
686 else
687 return NULL;
688 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689 if (PyList_Size(lz->saved) == 0)
690 return NULL;
691 it = PyObject_GetIter(lz->saved);
692 if (it == NULL)
693 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000694 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695 lz->it = it;
696 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000697 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000698 }
699}
700
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000701PyDoc_STRVAR(cycle_doc,
702"cycle(iterable) --> cycle object\n\
703\n\
704Return elements from the iterable until it is exhausted.\n\
705Then repeat the sequence indefinitely.");
706
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000707static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000708 PyObject_HEAD_INIT(NULL)
709 0, /* ob_size */
710 "itertools.cycle", /* tp_name */
711 sizeof(cycleobject), /* tp_basicsize */
712 0, /* tp_itemsize */
713 /* methods */
714 (destructor)cycle_dealloc, /* tp_dealloc */
715 0, /* tp_print */
716 0, /* tp_getattr */
717 0, /* tp_setattr */
718 0, /* tp_compare */
719 0, /* tp_repr */
720 0, /* tp_as_number */
721 0, /* tp_as_sequence */
722 0, /* tp_as_mapping */
723 0, /* tp_hash */
724 0, /* tp_call */
725 0, /* tp_str */
726 PyObject_GenericGetAttr, /* tp_getattro */
727 0, /* tp_setattro */
728 0, /* tp_as_buffer */
729 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
730 Py_TPFLAGS_BASETYPE, /* tp_flags */
731 cycle_doc, /* tp_doc */
732 (traverseproc)cycle_traverse, /* tp_traverse */
733 0, /* tp_clear */
734 0, /* tp_richcompare */
735 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000736 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000737 (iternextfunc)cycle_next, /* tp_iternext */
738 0, /* tp_methods */
739 0, /* tp_members */
740 0, /* tp_getset */
741 0, /* tp_base */
742 0, /* tp_dict */
743 0, /* tp_descr_get */
744 0, /* tp_descr_set */
745 0, /* tp_dictoffset */
746 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000747 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000748 cycle_new, /* tp_new */
749 PyObject_GC_Del, /* tp_free */
750};
751
752
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000753/* dropwhile object **********************************************************/
754
755typedef struct {
756 PyObject_HEAD
757 PyObject *func;
758 PyObject *it;
759 long start;
760} dropwhileobject;
761
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000762static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000763
764static PyObject *
765dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
766{
767 PyObject *func, *seq;
768 PyObject *it;
769 dropwhileobject *lz;
770
Georg Brandl02c42872005-08-26 06:42:30 +0000771 if (!_PyArg_NoKeywords("dropwhile()", kwds))
772 return NULL;
773
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000774 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
775 return NULL;
776
777 /* Get iterator. */
778 it = PyObject_GetIter(seq);
779 if (it == NULL)
780 return NULL;
781
782 /* create dropwhileobject structure */
783 lz = (dropwhileobject *)type->tp_alloc(type, 0);
784 if (lz == NULL) {
785 Py_DECREF(it);
786 return NULL;
787 }
788 Py_INCREF(func);
789 lz->func = func;
790 lz->it = it;
791 lz->start = 0;
792
793 return (PyObject *)lz;
794}
795
796static void
797dropwhile_dealloc(dropwhileobject *lz)
798{
799 PyObject_GC_UnTrack(lz);
800 Py_XDECREF(lz->func);
801 Py_XDECREF(lz->it);
802 lz->ob_type->tp_free(lz);
803}
804
805static int
806dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
807{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000808 Py_VISIT(lz->it);
809 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000810 return 0;
811}
812
813static PyObject *
814dropwhile_next(dropwhileobject *lz)
815{
816 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000817 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000818 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000819 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000820
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000821 assert(PyIter_Check(it));
822 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000823 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000824 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000825 if (item == NULL)
826 return NULL;
827 if (lz->start == 1)
828 return item;
829
830 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
831 if (good == NULL) {
832 Py_DECREF(item);
833 return NULL;
834 }
835 ok = PyObject_IsTrue(good);
836 Py_DECREF(good);
837 if (!ok) {
838 lz->start = 1;
839 return item;
840 }
841 Py_DECREF(item);
842 }
843}
844
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000845PyDoc_STRVAR(dropwhile_doc,
846"dropwhile(predicate, iterable) --> dropwhile object\n\
847\n\
848Drop items from the iterable while predicate(item) is true.\n\
849Afterwards, return every element until the iterable is exhausted.");
850
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000851static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000852 PyObject_HEAD_INIT(NULL)
853 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000854 "itertools.dropwhile", /* tp_name */
855 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000856 0, /* tp_itemsize */
857 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000858 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000859 0, /* tp_print */
860 0, /* tp_getattr */
861 0, /* tp_setattr */
862 0, /* tp_compare */
863 0, /* tp_repr */
864 0, /* tp_as_number */
865 0, /* tp_as_sequence */
866 0, /* tp_as_mapping */
867 0, /* tp_hash */
868 0, /* tp_call */
869 0, /* tp_str */
870 PyObject_GenericGetAttr, /* tp_getattro */
871 0, /* tp_setattro */
872 0, /* tp_as_buffer */
873 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
874 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000875 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876 (traverseproc)dropwhile_traverse, /* tp_traverse */
877 0, /* tp_clear */
878 0, /* tp_richcompare */
879 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000880 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000881 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000882 0, /* tp_methods */
883 0, /* tp_members */
884 0, /* tp_getset */
885 0, /* tp_base */
886 0, /* tp_dict */
887 0, /* tp_descr_get */
888 0, /* tp_descr_set */
889 0, /* tp_dictoffset */
890 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000891 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000892 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000893 PyObject_GC_Del, /* tp_free */
894};
895
896
897/* takewhile object **********************************************************/
898
899typedef struct {
900 PyObject_HEAD
901 PyObject *func;
902 PyObject *it;
903 long stop;
904} takewhileobject;
905
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000906static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000907
908static PyObject *
909takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
910{
911 PyObject *func, *seq;
912 PyObject *it;
913 takewhileobject *lz;
914
Georg Brandl02c42872005-08-26 06:42:30 +0000915 if (!_PyArg_NoKeywords("takewhile()", kwds))
916 return NULL;
917
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000918 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
919 return NULL;
920
921 /* Get iterator. */
922 it = PyObject_GetIter(seq);
923 if (it == NULL)
924 return NULL;
925
926 /* create takewhileobject structure */
927 lz = (takewhileobject *)type->tp_alloc(type, 0);
928 if (lz == NULL) {
929 Py_DECREF(it);
930 return NULL;
931 }
932 Py_INCREF(func);
933 lz->func = func;
934 lz->it = it;
935 lz->stop = 0;
936
937 return (PyObject *)lz;
938}
939
940static void
941takewhile_dealloc(takewhileobject *lz)
942{
943 PyObject_GC_UnTrack(lz);
944 Py_XDECREF(lz->func);
945 Py_XDECREF(lz->it);
946 lz->ob_type->tp_free(lz);
947}
948
949static int
950takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
951{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000952 Py_VISIT(lz->it);
953 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000954 return 0;
955}
956
957static PyObject *
958takewhile_next(takewhileobject *lz)
959{
960 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000961 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000962 long ok;
963
964 if (lz->stop == 1)
965 return NULL;
966
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000967 assert(PyIter_Check(it));
968 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000969 if (item == NULL)
970 return NULL;
971
972 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
973 if (good == NULL) {
974 Py_DECREF(item);
975 return NULL;
976 }
977 ok = PyObject_IsTrue(good);
978 Py_DECREF(good);
979 if (ok)
980 return item;
981 Py_DECREF(item);
982 lz->stop = 1;
983 return NULL;
984}
985
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000986PyDoc_STRVAR(takewhile_doc,
987"takewhile(predicate, iterable) --> takewhile object\n\
988\n\
989Return successive entries from an iterable as long as the \n\
990predicate evaluates to true for each entry.");
991
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000992static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993 PyObject_HEAD_INIT(NULL)
994 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000995 "itertools.takewhile", /* tp_name */
996 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000997 0, /* tp_itemsize */
998 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000999 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001000 0, /* tp_print */
1001 0, /* tp_getattr */
1002 0, /* tp_setattr */
1003 0, /* tp_compare */
1004 0, /* tp_repr */
1005 0, /* tp_as_number */
1006 0, /* tp_as_sequence */
1007 0, /* tp_as_mapping */
1008 0, /* tp_hash */
1009 0, /* tp_call */
1010 0, /* tp_str */
1011 PyObject_GenericGetAttr, /* tp_getattro */
1012 0, /* tp_setattro */
1013 0, /* tp_as_buffer */
1014 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1015 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001016 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001017 (traverseproc)takewhile_traverse, /* tp_traverse */
1018 0, /* tp_clear */
1019 0, /* tp_richcompare */
1020 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001021 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001022 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023 0, /* tp_methods */
1024 0, /* tp_members */
1025 0, /* tp_getset */
1026 0, /* tp_base */
1027 0, /* tp_dict */
1028 0, /* tp_descr_get */
1029 0, /* tp_descr_set */
1030 0, /* tp_dictoffset */
1031 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001032 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001033 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001034 PyObject_GC_Del, /* tp_free */
1035};
1036
1037
1038/* islice object ************************************************************/
1039
1040typedef struct {
1041 PyObject_HEAD
1042 PyObject *it;
1043 long next;
1044 long stop;
1045 long step;
1046 long cnt;
1047} isliceobject;
1048
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001049static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001050
1051static PyObject *
1052islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1053{
1054 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001055 long start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001056 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001057 int numargs;
1058 isliceobject *lz;
1059
Georg Brandl02c42872005-08-26 06:42:30 +00001060 if (!_PyArg_NoKeywords("islice()", kwds))
1061 return NULL;
1062
Raymond Hettingerb2594052004-12-05 09:25:51 +00001063 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001064 return NULL;
1065
Raymond Hettingerb2594052004-12-05 09:25:51 +00001066 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001067 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001068 if (a1 != Py_None) {
1069 stop = PyInt_AsLong(a1);
1070 if (stop == -1) {
1071 if (PyErr_Occurred())
1072 PyErr_Clear();
1073 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001074 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001075 return NULL;
1076 }
1077 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001078 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001079 if (a1 != Py_None)
1080 start = PyInt_AsLong(a1);
1081 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001082 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001083 if (a2 != Py_None) {
1084 stop = PyInt_AsLong(a2);
1085 if (stop == -1) {
1086 if (PyErr_Occurred())
1087 PyErr_Clear();
1088 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001089 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001090 return NULL;
1091 }
1092 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001093 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001094 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001095 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001096 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001097 return NULL;
1098 }
1099
Raymond Hettingerb2594052004-12-05 09:25:51 +00001100 if (a3 != NULL) {
1101 if (a3 != Py_None)
1102 step = PyInt_AsLong(a3);
1103 if (step == -1 && PyErr_Occurred())
1104 PyErr_Clear();
1105 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001106 if (step<1) {
1107 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001108 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001109 return NULL;
1110 }
1111
1112 /* Get iterator. */
1113 it = PyObject_GetIter(seq);
1114 if (it == NULL)
1115 return NULL;
1116
1117 /* create isliceobject structure */
1118 lz = (isliceobject *)type->tp_alloc(type, 0);
1119 if (lz == NULL) {
1120 Py_DECREF(it);
1121 return NULL;
1122 }
1123 lz->it = it;
1124 lz->next = start;
1125 lz->stop = stop;
1126 lz->step = step;
1127 lz->cnt = 0L;
1128
1129 return (PyObject *)lz;
1130}
1131
1132static void
1133islice_dealloc(isliceobject *lz)
1134{
1135 PyObject_GC_UnTrack(lz);
1136 Py_XDECREF(lz->it);
1137 lz->ob_type->tp_free(lz);
1138}
1139
1140static int
1141islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1142{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001143 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001144 return 0;
1145}
1146
1147static PyObject *
1148islice_next(isliceobject *lz)
1149{
1150 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001151 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001152 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001153 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001155 assert(PyIter_Check(it));
1156 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001157 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001158 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 if (item == NULL)
1160 return NULL;
1161 Py_DECREF(item);
1162 lz->cnt++;
1163 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001164 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001166 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001167 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001168 if (item == NULL)
1169 return NULL;
1170 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001171 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001172 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001173 if (lz->next < oldnext) /* Check for overflow */
1174 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 return item;
1176}
1177
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001178PyDoc_STRVAR(islice_doc,
1179"islice(iterable, [start,] stop [, step]) --> islice object\n\
1180\n\
1181Return an iterator whose next() method returns selected values from an\n\
1182iterable. If start is specified, will skip all preceding elements;\n\
1183otherwise, start defaults to zero. Step defaults to one. If\n\
1184specified as another value, step determines how many values are \n\
1185skipped between successive calls. Works like a slice() on a list\n\
1186but returns an iterator.");
1187
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001188static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001189 PyObject_HEAD_INIT(NULL)
1190 0, /* ob_size */
1191 "itertools.islice", /* tp_name */
1192 sizeof(isliceobject), /* tp_basicsize */
1193 0, /* tp_itemsize */
1194 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001195 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196 0, /* tp_print */
1197 0, /* tp_getattr */
1198 0, /* tp_setattr */
1199 0, /* tp_compare */
1200 0, /* tp_repr */
1201 0, /* tp_as_number */
1202 0, /* tp_as_sequence */
1203 0, /* tp_as_mapping */
1204 0, /* tp_hash */
1205 0, /* tp_call */
1206 0, /* tp_str */
1207 PyObject_GenericGetAttr, /* tp_getattro */
1208 0, /* tp_setattro */
1209 0, /* tp_as_buffer */
1210 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1211 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001212 islice_doc, /* tp_doc */
1213 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001214 0, /* tp_clear */
1215 0, /* tp_richcompare */
1216 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001217 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001218 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001219 0, /* tp_methods */
1220 0, /* tp_members */
1221 0, /* tp_getset */
1222 0, /* tp_base */
1223 0, /* tp_dict */
1224 0, /* tp_descr_get */
1225 0, /* tp_descr_set */
1226 0, /* tp_dictoffset */
1227 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001228 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001229 islice_new, /* tp_new */
1230 PyObject_GC_Del, /* tp_free */
1231};
1232
1233
1234/* starmap object ************************************************************/
1235
1236typedef struct {
1237 PyObject_HEAD
1238 PyObject *func;
1239 PyObject *it;
1240} starmapobject;
1241
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001242static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001243
1244static PyObject *
1245starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1246{
1247 PyObject *func, *seq;
1248 PyObject *it;
1249 starmapobject *lz;
1250
Georg Brandl02c42872005-08-26 06:42:30 +00001251 if (!_PyArg_NoKeywords("starmap()", kwds))
1252 return NULL;
1253
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1255 return NULL;
1256
1257 /* Get iterator. */
1258 it = PyObject_GetIter(seq);
1259 if (it == NULL)
1260 return NULL;
1261
1262 /* create starmapobject structure */
1263 lz = (starmapobject *)type->tp_alloc(type, 0);
1264 if (lz == NULL) {
1265 Py_DECREF(it);
1266 return NULL;
1267 }
1268 Py_INCREF(func);
1269 lz->func = func;
1270 lz->it = it;
1271
1272 return (PyObject *)lz;
1273}
1274
1275static void
1276starmap_dealloc(starmapobject *lz)
1277{
1278 PyObject_GC_UnTrack(lz);
1279 Py_XDECREF(lz->func);
1280 Py_XDECREF(lz->it);
1281 lz->ob_type->tp_free(lz);
1282}
1283
1284static int
1285starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1286{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001287 Py_VISIT(lz->it);
1288 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001289 return 0;
1290}
1291
1292static PyObject *
1293starmap_next(starmapobject *lz)
1294{
1295 PyObject *args;
1296 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001297 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001298
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001299 assert(PyIter_Check(it));
1300 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001301 if (args == NULL)
1302 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001303 if (!PyTuple_CheckExact(args)) {
1304 Py_DECREF(args);
1305 PyErr_SetString(PyExc_TypeError,
1306 "iterator must return a tuple");
1307 return NULL;
1308 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309 result = PyObject_Call(lz->func, args, NULL);
1310 Py_DECREF(args);
1311 return result;
1312}
1313
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314PyDoc_STRVAR(starmap_doc,
1315"starmap(function, sequence) --> starmap object\n\
1316\n\
1317Return an iterator whose values are returned from the function evaluated\n\
1318with a argument tuple taken from the given sequence.");
1319
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001320static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321 PyObject_HEAD_INIT(NULL)
1322 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001323 "itertools.starmap", /* tp_name */
1324 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001325 0, /* tp_itemsize */
1326 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001327 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328 0, /* tp_print */
1329 0, /* tp_getattr */
1330 0, /* tp_setattr */
1331 0, /* tp_compare */
1332 0, /* tp_repr */
1333 0, /* tp_as_number */
1334 0, /* tp_as_sequence */
1335 0, /* tp_as_mapping */
1336 0, /* tp_hash */
1337 0, /* tp_call */
1338 0, /* tp_str */
1339 PyObject_GenericGetAttr, /* tp_getattro */
1340 0, /* tp_setattro */
1341 0, /* tp_as_buffer */
1342 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1343 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001344 starmap_doc, /* tp_doc */
1345 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346 0, /* tp_clear */
1347 0, /* tp_richcompare */
1348 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001349 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001350 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351 0, /* tp_methods */
1352 0, /* tp_members */
1353 0, /* tp_getset */
1354 0, /* tp_base */
1355 0, /* tp_dict */
1356 0, /* tp_descr_get */
1357 0, /* tp_descr_set */
1358 0, /* tp_dictoffset */
1359 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001360 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001361 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362 PyObject_GC_Del, /* tp_free */
1363};
1364
1365
1366/* imap object ************************************************************/
1367
1368typedef struct {
1369 PyObject_HEAD
1370 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001371 PyObject *func;
1372} imapobject;
1373
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001374static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375
1376static PyObject *
1377imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1378{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001379 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380 imapobject *lz;
1381 int numargs, i;
1382
Georg Brandl02c42872005-08-26 06:42:30 +00001383 if (!_PyArg_NoKeywords("imap()", kwds))
1384 return NULL;
1385
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386 numargs = PyTuple_Size(args);
1387 if (numargs < 2) {
1388 PyErr_SetString(PyExc_TypeError,
1389 "imap() must have at least two arguments.");
1390 return NULL;
1391 }
1392
1393 iters = PyTuple_New(numargs-1);
1394 if (iters == NULL)
1395 return NULL;
1396
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001397 for (i=1 ; i<numargs ; i++) {
1398 /* Get iterator. */
1399 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1400 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001401 Py_DECREF(iters);
1402 return NULL;
1403 }
1404 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001405 }
1406
1407 /* create imapobject structure */
1408 lz = (imapobject *)type->tp_alloc(type, 0);
1409 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001410 Py_DECREF(iters);
1411 return NULL;
1412 }
1413 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001414 func = PyTuple_GET_ITEM(args, 0);
1415 Py_INCREF(func);
1416 lz->func = func;
1417
1418 return (PyObject *)lz;
1419}
1420
1421static void
1422imap_dealloc(imapobject *lz)
1423{
1424 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001425 Py_XDECREF(lz->iters);
1426 Py_XDECREF(lz->func);
1427 lz->ob_type->tp_free(lz);
1428}
1429
1430static int
1431imap_traverse(imapobject *lz, visitproc visit, void *arg)
1432{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001433 Py_VISIT(lz->iters);
1434 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435 return 0;
1436}
1437
Raymond Hettinger2012f172003-02-07 05:32:58 +00001438/*
1439imap() is an iterator version of __builtins__.map() except that it does
1440not have the None fill-in feature. That was intentionally left out for
1441the following reasons:
1442
1443 1) Itertools are designed to be easily combined and chained together.
1444 Having all tools stop with the shortest input is a unifying principle
1445 that makes it easier to combine finite iterators (supplying data) with
1446 infinite iterators like count() and repeat() (for supplying sequential
1447 or constant arguments to a function).
1448
1449 2) In typical use cases for combining itertools, having one finite data
1450 supplier run out before another is likely to be an error condition which
1451 should not pass silently by automatically supplying None.
1452
1453 3) The use cases for automatic None fill-in are rare -- not many functions
1454 do something useful when a parameter suddenly switches type and becomes
1455 None.
1456
1457 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001458 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001459
1460 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1461*/
1462
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463static PyObject *
1464imap_next(imapobject *lz)
1465{
1466 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001467 PyObject *argtuple;
1468 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469 int numargs, i;
1470
1471 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001472 argtuple = PyTuple_New(numargs);
1473 if (argtuple == NULL)
1474 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001476 for (i=0 ; i<numargs ; i++) {
1477 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1478 if (val == NULL) {
1479 Py_DECREF(argtuple);
1480 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001482 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001484 if (lz->func == Py_None)
1485 return argtuple;
1486 result = PyObject_Call(lz->func, argtuple, NULL);
1487 Py_DECREF(argtuple);
1488 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001489}
1490
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001491PyDoc_STRVAR(imap_doc,
1492"imap(func, *iterables) --> imap object\n\
1493\n\
1494Make an iterator that computes the function using arguments from\n\
1495each of the iterables. Like map() except that it returns\n\
1496an iterator instead of a list and that it stops when the shortest\n\
1497iterable is exhausted instead of filling in None for shorter\n\
1498iterables.");
1499
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001500static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501 PyObject_HEAD_INIT(NULL)
1502 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001503 "itertools.imap", /* tp_name */
1504 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505 0, /* tp_itemsize */
1506 /* methods */
1507 (destructor)imap_dealloc, /* tp_dealloc */
1508 0, /* tp_print */
1509 0, /* tp_getattr */
1510 0, /* tp_setattr */
1511 0, /* tp_compare */
1512 0, /* tp_repr */
1513 0, /* tp_as_number */
1514 0, /* tp_as_sequence */
1515 0, /* tp_as_mapping */
1516 0, /* tp_hash */
1517 0, /* tp_call */
1518 0, /* tp_str */
1519 PyObject_GenericGetAttr, /* tp_getattro */
1520 0, /* tp_setattro */
1521 0, /* tp_as_buffer */
1522 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1523 Py_TPFLAGS_BASETYPE, /* tp_flags */
1524 imap_doc, /* tp_doc */
1525 (traverseproc)imap_traverse, /* tp_traverse */
1526 0, /* tp_clear */
1527 0, /* tp_richcompare */
1528 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001529 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530 (iternextfunc)imap_next, /* tp_iternext */
1531 0, /* tp_methods */
1532 0, /* tp_members */
1533 0, /* tp_getset */
1534 0, /* tp_base */
1535 0, /* tp_dict */
1536 0, /* tp_descr_get */
1537 0, /* tp_descr_set */
1538 0, /* tp_dictoffset */
1539 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001540 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001541 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001542 PyObject_GC_Del, /* tp_free */
1543};
1544
1545
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001546/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001547
1548typedef struct {
1549 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001550 long tuplesize;
1551 long iternum; /* which iterator is active */
1552 PyObject *ittuple; /* tuple of iterators */
1553} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001555static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556
1557static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001558chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001560 chainobject *lz;
1561 int tuplesize = PySequence_Length(args);
1562 int i;
1563 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564
Georg Brandl02c42872005-08-26 06:42:30 +00001565 if (!_PyArg_NoKeywords("chain()", kwds))
1566 return NULL;
1567
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001568 /* obtain iterators */
1569 assert(PyTuple_Check(args));
1570 ittuple = PyTuple_New(tuplesize);
1571 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001572 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001573 for (i=0; i < tuplesize; ++i) {
1574 PyObject *item = PyTuple_GET_ITEM(args, i);
1575 PyObject *it = PyObject_GetIter(item);
1576 if (it == NULL) {
1577 if (PyErr_ExceptionMatches(PyExc_TypeError))
1578 PyErr_Format(PyExc_TypeError,
1579 "chain argument #%d must support iteration",
1580 i+1);
1581 Py_DECREF(ittuple);
1582 return NULL;
1583 }
1584 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585 }
1586
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001587 /* create chainobject structure */
1588 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001589 if (lz == NULL) {
1590 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001591 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001592 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001593
1594 lz->ittuple = ittuple;
1595 lz->iternum = 0;
1596 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001597
1598 return (PyObject *)lz;
1599}
1600
1601static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001602chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001603{
1604 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001605 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606 lz->ob_type->tp_free(lz);
1607}
1608
Raymond Hettinger2012f172003-02-07 05:32:58 +00001609static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001610chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001611{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001612 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001613 return 0;
1614}
1615
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001617chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001619 PyObject *it;
1620 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001622 while (lz->iternum < lz->tuplesize) {
1623 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1624 item = PyIter_Next(it);
1625 if (item != NULL)
1626 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001627 if (PyErr_Occurred()) {
1628 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1629 PyErr_Clear();
1630 else
1631 return NULL;
1632 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001633 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001634 }
1635 return NULL;
1636}
1637
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001638PyDoc_STRVAR(chain_doc,
1639"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001641Return a chain object whose .next() method returns elements from the\n\
1642first iterable until it is exhausted, then elements from the next\n\
1643iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001645static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001646 PyObject_HEAD_INIT(NULL)
1647 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001648 "itertools.chain", /* tp_name */
1649 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650 0, /* tp_itemsize */
1651 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001652 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001653 0, /* tp_print */
1654 0, /* tp_getattr */
1655 0, /* tp_setattr */
1656 0, /* tp_compare */
1657 0, /* tp_repr */
1658 0, /* tp_as_number */
1659 0, /* tp_as_sequence */
1660 0, /* tp_as_mapping */
1661 0, /* tp_hash */
1662 0, /* tp_call */
1663 0, /* tp_str */
1664 PyObject_GenericGetAttr, /* tp_getattro */
1665 0, /* tp_setattro */
1666 0, /* tp_as_buffer */
1667 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1668 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001669 chain_doc, /* tp_doc */
1670 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671 0, /* tp_clear */
1672 0, /* tp_richcompare */
1673 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001674 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001675 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676 0, /* tp_methods */
1677 0, /* tp_members */
1678 0, /* tp_getset */
1679 0, /* tp_base */
1680 0, /* tp_dict */
1681 0, /* tp_descr_get */
1682 0, /* tp_descr_set */
1683 0, /* tp_dictoffset */
1684 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001685 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001686 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687 PyObject_GC_Del, /* tp_free */
1688};
1689
1690
1691/* ifilter object ************************************************************/
1692
1693typedef struct {
1694 PyObject_HEAD
1695 PyObject *func;
1696 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697} ifilterobject;
1698
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001699static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001700
1701static PyObject *
1702ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1703{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001704 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705 PyObject *it;
1706 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001707
Georg Brandl02c42872005-08-26 06:42:30 +00001708 if (!_PyArg_NoKeywords("ifilter()", kwds))
1709 return NULL;
1710
Raymond Hettinger60eca932003-02-09 06:40:58 +00001711 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712 return NULL;
1713
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001714 /* Get iterator. */
1715 it = PyObject_GetIter(seq);
1716 if (it == NULL)
1717 return NULL;
1718
1719 /* create ifilterobject structure */
1720 lz = (ifilterobject *)type->tp_alloc(type, 0);
1721 if (lz == NULL) {
1722 Py_DECREF(it);
1723 return NULL;
1724 }
1725 Py_INCREF(func);
1726 lz->func = func;
1727 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001728
1729 return (PyObject *)lz;
1730}
1731
1732static void
1733ifilter_dealloc(ifilterobject *lz)
1734{
1735 PyObject_GC_UnTrack(lz);
1736 Py_XDECREF(lz->func);
1737 Py_XDECREF(lz->it);
1738 lz->ob_type->tp_free(lz);
1739}
1740
1741static int
1742ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1743{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001744 Py_VISIT(lz->it);
1745 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746 return 0;
1747}
1748
1749static PyObject *
1750ifilter_next(ifilterobject *lz)
1751{
1752 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001753 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001754 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001755 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001756
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001757 assert(PyIter_Check(it));
1758 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001760 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001761 if (item == NULL)
1762 return NULL;
1763
1764 if (lz->func == Py_None) {
1765 ok = PyObject_IsTrue(item);
1766 } else {
1767 PyObject *good;
1768 good = PyObject_CallFunctionObjArgs(lz->func,
1769 item, NULL);
1770 if (good == NULL) {
1771 Py_DECREF(item);
1772 return NULL;
1773 }
1774 ok = PyObject_IsTrue(good);
1775 Py_DECREF(good);
1776 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001777 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001778 return item;
1779 Py_DECREF(item);
1780 }
1781}
1782
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001783PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001784"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001785\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001786Return those items of sequence for which function(item) is true.\n\
1787If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001788
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001789static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001790 PyObject_HEAD_INIT(NULL)
1791 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001792 "itertools.ifilter", /* tp_name */
1793 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001794 0, /* tp_itemsize */
1795 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001796 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001797 0, /* tp_print */
1798 0, /* tp_getattr */
1799 0, /* tp_setattr */
1800 0, /* tp_compare */
1801 0, /* tp_repr */
1802 0, /* tp_as_number */
1803 0, /* tp_as_sequence */
1804 0, /* tp_as_mapping */
1805 0, /* tp_hash */
1806 0, /* tp_call */
1807 0, /* tp_str */
1808 PyObject_GenericGetAttr, /* tp_getattro */
1809 0, /* tp_setattro */
1810 0, /* tp_as_buffer */
1811 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1812 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001813 ifilter_doc, /* tp_doc */
1814 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001815 0, /* tp_clear */
1816 0, /* tp_richcompare */
1817 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001818 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001819 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001820 0, /* tp_methods */
1821 0, /* tp_members */
1822 0, /* tp_getset */
1823 0, /* tp_base */
1824 0, /* tp_dict */
1825 0, /* tp_descr_get */
1826 0, /* tp_descr_set */
1827 0, /* tp_dictoffset */
1828 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001829 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001830 ifilter_new, /* tp_new */
1831 PyObject_GC_Del, /* tp_free */
1832};
1833
1834
1835/* ifilterfalse object ************************************************************/
1836
1837typedef struct {
1838 PyObject_HEAD
1839 PyObject *func;
1840 PyObject *it;
1841} ifilterfalseobject;
1842
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001843static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001844
1845static PyObject *
1846ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1847{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001848 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001849 PyObject *it;
1850 ifilterfalseobject *lz;
1851
Georg Brandl02c42872005-08-26 06:42:30 +00001852 if (!_PyArg_NoKeywords("ifilterfalse()", kwds))
1853 return NULL;
1854
Raymond Hettinger60eca932003-02-09 06:40:58 +00001855 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1856 return NULL;
1857
1858 /* Get iterator. */
1859 it = PyObject_GetIter(seq);
1860 if (it == NULL)
1861 return NULL;
1862
1863 /* create ifilterfalseobject structure */
1864 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1865 if (lz == NULL) {
1866 Py_DECREF(it);
1867 return NULL;
1868 }
1869 Py_INCREF(func);
1870 lz->func = func;
1871 lz->it = it;
1872
1873 return (PyObject *)lz;
1874}
1875
1876static void
1877ifilterfalse_dealloc(ifilterfalseobject *lz)
1878{
1879 PyObject_GC_UnTrack(lz);
1880 Py_XDECREF(lz->func);
1881 Py_XDECREF(lz->it);
1882 lz->ob_type->tp_free(lz);
1883}
1884
1885static int
1886ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1887{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001888 Py_VISIT(lz->it);
1889 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001890 return 0;
1891}
1892
1893static PyObject *
1894ifilterfalse_next(ifilterfalseobject *lz)
1895{
1896 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001897 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001898 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001899 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001900
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001901 assert(PyIter_Check(it));
1902 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001903 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001904 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001905 if (item == NULL)
1906 return NULL;
1907
1908 if (lz->func == Py_None) {
1909 ok = PyObject_IsTrue(item);
1910 } else {
1911 PyObject *good;
1912 good = PyObject_CallFunctionObjArgs(lz->func,
1913 item, NULL);
1914 if (good == NULL) {
1915 Py_DECREF(item);
1916 return NULL;
1917 }
1918 ok = PyObject_IsTrue(good);
1919 Py_DECREF(good);
1920 }
1921 if (!ok)
1922 return item;
1923 Py_DECREF(item);
1924 }
1925}
1926
Raymond Hettinger60eca932003-02-09 06:40:58 +00001927PyDoc_STRVAR(ifilterfalse_doc,
1928"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1929\n\
1930Return those items of sequence for which function(item) is false.\n\
1931If function is None, return the items that are false.");
1932
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001933static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001934 PyObject_HEAD_INIT(NULL)
1935 0, /* ob_size */
1936 "itertools.ifilterfalse", /* tp_name */
1937 sizeof(ifilterfalseobject), /* tp_basicsize */
1938 0, /* tp_itemsize */
1939 /* methods */
1940 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1941 0, /* tp_print */
1942 0, /* tp_getattr */
1943 0, /* tp_setattr */
1944 0, /* tp_compare */
1945 0, /* tp_repr */
1946 0, /* tp_as_number */
1947 0, /* tp_as_sequence */
1948 0, /* tp_as_mapping */
1949 0, /* tp_hash */
1950 0, /* tp_call */
1951 0, /* tp_str */
1952 PyObject_GenericGetAttr, /* tp_getattro */
1953 0, /* tp_setattro */
1954 0, /* tp_as_buffer */
1955 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1956 Py_TPFLAGS_BASETYPE, /* tp_flags */
1957 ifilterfalse_doc, /* tp_doc */
1958 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1959 0, /* tp_clear */
1960 0, /* tp_richcompare */
1961 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001962 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001963 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1964 0, /* tp_methods */
1965 0, /* tp_members */
1966 0, /* tp_getset */
1967 0, /* tp_base */
1968 0, /* tp_dict */
1969 0, /* tp_descr_get */
1970 0, /* tp_descr_set */
1971 0, /* tp_dictoffset */
1972 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001973 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001974 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001975 PyObject_GC_Del, /* tp_free */
1976};
1977
1978
1979/* count object ************************************************************/
1980
1981typedef struct {
1982 PyObject_HEAD
1983 long cnt;
1984} countobject;
1985
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001986static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001987
1988static PyObject *
1989count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1990{
1991 countobject *lz;
1992 long cnt = 0;
1993
Georg Brandl02c42872005-08-26 06:42:30 +00001994 if (!_PyArg_NoKeywords("count()", kwds))
1995 return NULL;
1996
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001997 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1998 return NULL;
1999
2000 /* create countobject structure */
2001 lz = (countobject *)PyObject_New(countobject, &count_type);
2002 if (lz == NULL)
2003 return NULL;
2004 lz->cnt = cnt;
2005
2006 return (PyObject *)lz;
2007}
2008
2009static PyObject *
2010count_next(countobject *lz)
2011{
2012 return PyInt_FromLong(lz->cnt++);
2013}
2014
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002015static PyObject *
2016count_repr(countobject *lz)
2017{
Brett Cannon00468392004-04-13 02:43:53 +00002018 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002019}
2020
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002021PyDoc_STRVAR(count_doc,
2022"count([firstval]) --> count object\n\
2023\n\
2024Return a count object whose .next() method returns consecutive\n\
2025integers starting from zero or, if specified, from firstval.");
2026
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002027static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002028 PyObject_HEAD_INIT(NULL)
2029 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002030 "itertools.count", /* tp_name */
2031 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002032 0, /* tp_itemsize */
2033 /* methods */
2034 (destructor)PyObject_Del, /* tp_dealloc */
2035 0, /* tp_print */
2036 0, /* tp_getattr */
2037 0, /* tp_setattr */
2038 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002039 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002040 0, /* tp_as_number */
2041 0, /* tp_as_sequence */
2042 0, /* tp_as_mapping */
2043 0, /* tp_hash */
2044 0, /* tp_call */
2045 0, /* tp_str */
2046 PyObject_GenericGetAttr, /* tp_getattro */
2047 0, /* tp_setattro */
2048 0, /* tp_as_buffer */
2049 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002050 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002051 0, /* tp_traverse */
2052 0, /* tp_clear */
2053 0, /* tp_richcompare */
2054 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002055 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002056 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002057 0, /* tp_methods */
2058 0, /* tp_members */
2059 0, /* tp_getset */
2060 0, /* tp_base */
2061 0, /* tp_dict */
2062 0, /* tp_descr_get */
2063 0, /* tp_descr_set */
2064 0, /* tp_dictoffset */
2065 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002066 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002067 count_new, /* tp_new */
2068};
2069
2070
2071/* izip object ************************************************************/
2072
2073#include "Python.h"
2074
2075typedef struct {
2076 PyObject_HEAD
2077 long tuplesize;
2078 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002079 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002080} izipobject;
2081
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002082static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002083
2084static PyObject *
2085izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2086{
2087 izipobject *lz;
2088 int i;
2089 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002090 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002091 int tuplesize = PySequence_Length(args);
2092
Georg Brandl02c42872005-08-26 06:42:30 +00002093 if (!_PyArg_NoKeywords("izip()", kwds))
2094 return NULL;
2095
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002096 /* args must be a tuple */
2097 assert(PyTuple_Check(args));
2098
2099 /* obtain iterators */
2100 ittuple = PyTuple_New(tuplesize);
2101 if(ittuple == NULL)
2102 return NULL;
2103 for (i=0; i < tuplesize; ++i) {
2104 PyObject *item = PyTuple_GET_ITEM(args, i);
2105 PyObject *it = PyObject_GetIter(item);
2106 if (it == NULL) {
2107 if (PyErr_ExceptionMatches(PyExc_TypeError))
2108 PyErr_Format(PyExc_TypeError,
2109 "izip argument #%d must support iteration",
2110 i+1);
2111 Py_DECREF(ittuple);
2112 return NULL;
2113 }
2114 PyTuple_SET_ITEM(ittuple, i, it);
2115 }
2116
Raymond Hettinger2012f172003-02-07 05:32:58 +00002117 /* create a result holder */
2118 result = PyTuple_New(tuplesize);
2119 if (result == NULL) {
2120 Py_DECREF(ittuple);
2121 return NULL;
2122 }
2123 for (i=0 ; i < tuplesize ; i++) {
2124 Py_INCREF(Py_None);
2125 PyTuple_SET_ITEM(result, i, Py_None);
2126 }
2127
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002128 /* create izipobject structure */
2129 lz = (izipobject *)type->tp_alloc(type, 0);
2130 if (lz == NULL) {
2131 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002132 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002133 return NULL;
2134 }
2135 lz->ittuple = ittuple;
2136 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002137 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002138
2139 return (PyObject *)lz;
2140}
2141
2142static void
2143izip_dealloc(izipobject *lz)
2144{
2145 PyObject_GC_UnTrack(lz);
2146 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002147 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002148 lz->ob_type->tp_free(lz);
2149}
2150
2151static int
2152izip_traverse(izipobject *lz, visitproc visit, void *arg)
2153{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002154 Py_VISIT(lz->ittuple);
2155 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002156 return 0;
2157}
2158
2159static PyObject *
2160izip_next(izipobject *lz)
2161{
2162 int i;
2163 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002164 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002165 PyObject *it;
2166 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002167 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002168
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002169 if (tuplesize == 0)
2170 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002171 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002172 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002173 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002174 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002175 assert(PyIter_Check(it));
2176 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002177 if (item == NULL) {
2178 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002179 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002180 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002181 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002182 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002183 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002184 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002185 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002186 result = PyTuple_New(tuplesize);
2187 if (result == NULL)
2188 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002189 for (i=0 ; i < tuplesize ; i++) {
2190 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002191 assert(PyIter_Check(it));
2192 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002193 if (item == NULL) {
2194 Py_DECREF(result);
2195 return NULL;
2196 }
2197 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002198 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002199 }
2200 return result;
2201}
2202
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002203PyDoc_STRVAR(izip_doc,
2204"izip(iter1 [,iter2 [...]]) --> izip object\n\
2205\n\
2206Return a izip object whose .next() method returns a tuple where\n\
2207the i-th element comes from the i-th iterable argument. The .next()\n\
2208method continues until the shortest iterable in the argument sequence\n\
2209is exhausted and then it raises StopIteration. Works like the zip()\n\
2210function but consumes less memory by returning an iterator instead of\n\
2211a list.");
2212
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002213static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002214 PyObject_HEAD_INIT(NULL)
2215 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002216 "itertools.izip", /* tp_name */
2217 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002218 0, /* tp_itemsize */
2219 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002220 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002221 0, /* tp_print */
2222 0, /* tp_getattr */
2223 0, /* tp_setattr */
2224 0, /* tp_compare */
2225 0, /* tp_repr */
2226 0, /* tp_as_number */
2227 0, /* tp_as_sequence */
2228 0, /* tp_as_mapping */
2229 0, /* tp_hash */
2230 0, /* tp_call */
2231 0, /* tp_str */
2232 PyObject_GenericGetAttr, /* tp_getattro */
2233 0, /* tp_setattro */
2234 0, /* tp_as_buffer */
2235 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2236 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002237 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002238 (traverseproc)izip_traverse, /* tp_traverse */
2239 0, /* tp_clear */
2240 0, /* tp_richcompare */
2241 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002242 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002243 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002244 0, /* tp_methods */
2245 0, /* tp_members */
2246 0, /* tp_getset */
2247 0, /* tp_base */
2248 0, /* tp_dict */
2249 0, /* tp_descr_get */
2250 0, /* tp_descr_set */
2251 0, /* tp_dictoffset */
2252 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002253 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002254 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002255 PyObject_GC_Del, /* tp_free */
2256};
2257
2258
2259/* repeat object ************************************************************/
2260
2261typedef struct {
2262 PyObject_HEAD
2263 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002264 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002265} repeatobject;
2266
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002267static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002268
2269static PyObject *
2270repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2271{
2272 repeatobject *ro;
2273 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002274 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002275
Georg Brandl02c42872005-08-26 06:42:30 +00002276 if (!_PyArg_NoKeywords("repeat()", kwds))
2277 return NULL;
2278
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002279 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002280 return NULL;
2281
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002282 if (PyTuple_Size(args) == 2 && cnt < 0)
2283 cnt = 0;
2284
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002285 ro = (repeatobject *)type->tp_alloc(type, 0);
2286 if (ro == NULL)
2287 return NULL;
2288 Py_INCREF(element);
2289 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002290 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002291 return (PyObject *)ro;
2292}
2293
2294static void
2295repeat_dealloc(repeatobject *ro)
2296{
2297 PyObject_GC_UnTrack(ro);
2298 Py_XDECREF(ro->element);
2299 ro->ob_type->tp_free(ro);
2300}
2301
2302static int
2303repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2304{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002305 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002306 return 0;
2307}
2308
2309static PyObject *
2310repeat_next(repeatobject *ro)
2311{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002312 if (ro->cnt == 0)
2313 return NULL;
2314 if (ro->cnt > 0)
2315 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002316 Py_INCREF(ro->element);
2317 return ro->element;
2318}
2319
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002320static PyObject *
2321repeat_repr(repeatobject *ro)
2322{
2323 PyObject *result, *objrepr;
2324
2325 objrepr = PyObject_Repr(ro->element);
2326 if (objrepr == NULL)
2327 return NULL;
2328
2329 if (ro->cnt == -1)
2330 result = PyString_FromFormat("repeat(%s)",
2331 PyString_AS_STRING(objrepr));
2332 else
Brett Cannon00468392004-04-13 02:43:53 +00002333 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002334 PyString_AS_STRING(objrepr), ro->cnt);
2335 Py_DECREF(objrepr);
2336 return result;
2337}
2338
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002339static int
2340repeat_len(repeatobject *ro)
2341{
2342 if (ro->cnt == -1)
2343 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2344 return (int)(ro->cnt);
2345}
2346
2347static PySequenceMethods repeat_as_sequence = {
2348 (inquiry)repeat_len, /* sq_length */
2349 0, /* sq_concat */
2350};
2351
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002352PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002353"repeat(element [,times]) -> create an iterator which returns the element\n\
2354for the specified number of times. If not specified, returns the element\n\
2355endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002356
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002357static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002358 PyObject_HEAD_INIT(NULL)
2359 0, /* ob_size */
2360 "itertools.repeat", /* tp_name */
2361 sizeof(repeatobject), /* tp_basicsize */
2362 0, /* tp_itemsize */
2363 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002364 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002365 0, /* tp_print */
2366 0, /* tp_getattr */
2367 0, /* tp_setattr */
2368 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002369 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002370 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002371 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002372 0, /* tp_as_mapping */
2373 0, /* tp_hash */
2374 0, /* tp_call */
2375 0, /* tp_str */
2376 PyObject_GenericGetAttr, /* tp_getattro */
2377 0, /* tp_setattro */
2378 0, /* tp_as_buffer */
2379 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2380 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002381 repeat_doc, /* tp_doc */
2382 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002383 0, /* tp_clear */
2384 0, /* tp_richcompare */
2385 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002386 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002387 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002388 0, /* tp_methods */
2389 0, /* tp_members */
2390 0, /* tp_getset */
2391 0, /* tp_base */
2392 0, /* tp_dict */
2393 0, /* tp_descr_get */
2394 0, /* tp_descr_set */
2395 0, /* tp_dictoffset */
2396 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002397 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002398 repeat_new, /* tp_new */
2399 PyObject_GC_Del, /* tp_free */
2400};
2401
2402
2403/* module level code ********************************************************/
2404
2405PyDoc_STRVAR(module_doc,
2406"Functional tools for creating and using iterators.\n\
2407\n\
2408Infinite iterators:\n\
2409count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002410cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002411repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002412\n\
2413Iterators terminating on the shortest input sequence:\n\
2414izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002415ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2416ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002417islice(seq, [start,] stop [, step]) --> elements from\n\
2418 seq[start:stop:step]\n\
2419imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2420starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002421tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002422chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002423takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2424dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002425groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002426");
2427
2428
Raymond Hettingerad983e72003-11-12 14:32:26 +00002429static PyMethodDef module_methods[] = {
2430 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2431 {NULL, NULL} /* sentinel */
2432};
2433
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002434PyMODINIT_FUNC
2435inititertools(void)
2436{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002437 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002438 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002439 char *name;
2440 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002441 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002442 &dropwhile_type,
2443 &takewhile_type,
2444 &islice_type,
2445 &starmap_type,
2446 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002447 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002448 &ifilter_type,
2449 &ifilterfalse_type,
2450 &count_type,
2451 &izip_type,
2452 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002453 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002454 NULL
2455 };
2456
Skip Montanarof3938fd2004-02-10 20:27:40 +00002457 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002458 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002459
Raymond Hettinger60eca932003-02-09 06:40:58 +00002460 for (i=0 ; typelist[i] != NULL ; i++) {
2461 if (PyType_Ready(typelist[i]) < 0)
2462 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002463 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002464 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002465 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002466 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002467 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002468
2469 if (PyType_Ready(&teedataobject_type) < 0)
2470 return;
2471 if (PyType_Ready(&tee_type) < 0)
2472 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002473 if (PyType_Ready(&_grouper_type) < 0)
2474 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002475}