blob: 3445dccb4a764c33c38c852426c061932b474b3a [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
3
4/* Itertools module written and maintained
5 by Raymond D. Hettinger <python@rcn.com>
6 Copyright (c) 2003 Python Software Foundation.
7 All rights reserved.
8*/
9
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000010
11/* groupby object ***********************************************************/
12
13typedef struct {
14 PyObject_HEAD
15 PyObject *it;
16 PyObject *keyfunc;
17 PyObject *tgtkey;
18 PyObject *currkey;
19 PyObject *currvalue;
20} groupbyobject;
21
22static PyTypeObject groupby_type;
23static PyObject *_grouper_create(groupbyobject *, PyObject *);
24
25static PyObject *
26groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
27{
28 static char *kwargs[] = {"iterable", "key", NULL};
29 groupbyobject *gbo;
30 PyObject *it, *keyfunc = Py_None;
31
32 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
33 &it, &keyfunc))
34 return NULL;
35
36 gbo = (groupbyobject *)type->tp_alloc(type, 0);
37 if (gbo == NULL)
38 return NULL;
39 gbo->tgtkey = NULL;
40 gbo->currkey = NULL;
41 gbo->currvalue = NULL;
42 gbo->keyfunc = keyfunc;
43 Py_INCREF(keyfunc);
44 gbo->it = PyObject_GetIter(it);
45 if (gbo->it == NULL) {
46 Py_DECREF(gbo);
47 return NULL;
48 }
49 return (PyObject *)gbo;
50}
51
52static void
53groupby_dealloc(groupbyobject *gbo)
54{
55 PyObject_GC_UnTrack(gbo);
56 Py_XDECREF(gbo->it);
57 Py_XDECREF(gbo->keyfunc);
58 Py_XDECREF(gbo->tgtkey);
59 Py_XDECREF(gbo->currkey);
60 Py_XDECREF(gbo->currvalue);
61 gbo->ob_type->tp_free(gbo);
62}
63
64static int
65groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
66{
67 int err;
68
69 if (gbo->it) {
70 err = visit(gbo->it, arg);
71 if (err)
72 return err;
73 }
74 if (gbo->keyfunc) {
75 err = visit(gbo->keyfunc, arg);
76 if (err)
77 return err;
78 }
79 if (gbo->tgtkey) {
80 err = visit(gbo->tgtkey, arg);
81 if (err)
82 return err;
83 }
84 if (gbo->currkey) {
85 err = visit(gbo->currkey, arg);
86 if (err)
87 return err;
88 }
89 if (gbo->currvalue) {
90 err = visit(gbo->currvalue, arg);
91 if (err)
92 return err;
93 }
94 return 0;
95}
96
97static PyObject *
98groupby_next(groupbyobject *gbo)
99{
100 PyObject *newvalue, *newkey, *r, *grouper;
101
102 /* skip to next iteration group */
103 for (;;) {
104 if (gbo->currkey == NULL)
105 /* pass */;
106 else if (gbo->tgtkey == NULL)
107 break;
108 else {
109 int rcmp;
110
111 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
112 gbo->currkey, Py_EQ);
113 if (rcmp == -1)
114 return NULL;
115 else if (rcmp == 0)
116 break;
117 }
118
119 newvalue = PyIter_Next(gbo->it);
120 if (newvalue == NULL)
121 return NULL;
122
123 if (gbo->keyfunc == Py_None) {
124 newkey = newvalue;
125 Py_INCREF(newvalue);
126 } else {
127 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
128 newvalue, NULL);
129 if (newkey == NULL) {
130 Py_DECREF(newvalue);
131 return NULL;
132 }
133 }
134
135 Py_XDECREF(gbo->currkey);
136 gbo->currkey = newkey;
137 Py_XDECREF(gbo->currvalue);
138 gbo->currvalue = newvalue;
139 }
140
141 Py_XDECREF(gbo->tgtkey);
142 gbo->tgtkey = gbo->currkey;
143 Py_INCREF(gbo->currkey);
144
145 grouper = _grouper_create(gbo, gbo->tgtkey);
146 if (grouper == NULL)
147 return NULL;
148
149 r = PyTuple_Pack(2, gbo->currkey, grouper);
150 Py_DECREF(grouper);
151 return r;
152}
153
154PyDoc_STRVAR(groupby_doc,
155"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
156(key, sub-iterator) grouped by each value of key(value).\n");
157
158static PyTypeObject groupby_type = {
159 PyObject_HEAD_INIT(NULL)
160 0, /* ob_size */
161 "itertools.groupby", /* tp_name */
162 sizeof(groupbyobject), /* tp_basicsize */
163 0, /* tp_itemsize */
164 /* methods */
165 (destructor)groupby_dealloc, /* tp_dealloc */
166 0, /* tp_print */
167 0, /* tp_getattr */
168 0, /* tp_setattr */
169 0, /* tp_compare */
170 0, /* tp_repr */
171 0, /* tp_as_number */
172 0, /* tp_as_sequence */
173 0, /* tp_as_mapping */
174 0, /* tp_hash */
175 0, /* tp_call */
176 0, /* tp_str */
177 PyObject_GenericGetAttr, /* tp_getattro */
178 0, /* tp_setattro */
179 0, /* tp_as_buffer */
180 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
181 Py_TPFLAGS_BASETYPE, /* tp_flags */
182 groupby_doc, /* tp_doc */
183 (traverseproc)groupby_traverse, /* tp_traverse */
184 0, /* tp_clear */
185 0, /* tp_richcompare */
186 0, /* tp_weaklistoffset */
187 PyObject_SelfIter, /* tp_iter */
188 (iternextfunc)groupby_next, /* tp_iternext */
189 0, /* tp_methods */
190 0, /* tp_members */
191 0, /* tp_getset */
192 0, /* tp_base */
193 0, /* tp_dict */
194 0, /* tp_descr_get */
195 0, /* tp_descr_set */
196 0, /* tp_dictoffset */
197 0, /* tp_init */
198 0, /* tp_alloc */
199 groupby_new, /* tp_new */
200 PyObject_GC_Del, /* tp_free */
201};
202
203
204/* _grouper object (internal) ************************************************/
205
206typedef struct {
207 PyObject_HEAD
208 PyObject *parent;
209 PyObject *tgtkey;
210} _grouperobject;
211
212static PyTypeObject _grouper_type;
213
214static PyObject *
215_grouper_create(groupbyobject *parent, PyObject *tgtkey)
216{
217 _grouperobject *igo;
218
219 igo = PyObject_New(_grouperobject, &_grouper_type);
220 if (igo == NULL)
221 return NULL;
222 igo->parent = (PyObject *)parent;
223 Py_INCREF(parent);
224 igo->tgtkey = tgtkey;
225 Py_INCREF(tgtkey);
226
227 return (PyObject *)igo;
228}
229
230static void
231_grouper_dealloc(_grouperobject *igo)
232{
233 Py_DECREF(igo->parent);
234 Py_DECREF(igo->tgtkey);
235 PyObject_Del(igo);
236}
237
238static PyObject *
239_grouper_next(_grouperobject *igo)
240{
241 groupbyobject *gbo = (groupbyobject *)igo->parent;
242 PyObject *newvalue, *newkey, *r;
243 int rcmp;
244
245 if (gbo->currvalue == NULL) {
246 newvalue = PyIter_Next(gbo->it);
247 if (newvalue == NULL)
248 return NULL;
249
250 if (gbo->keyfunc == Py_None) {
251 newkey = newvalue;
252 Py_INCREF(newvalue);
253 } else {
254 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
255 newvalue, NULL);
256 if (newkey == NULL) {
257 Py_DECREF(newvalue);
258 return NULL;
259 }
260 }
261
262 assert(gbo->currkey == NULL);
263 gbo->currkey = newkey;
264 gbo->currvalue = newvalue;
265 }
266
267 assert(gbo->currkey != NULL);
268 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
269 if (rcmp <= 0)
270 /* got any error or current group is end */
271 return NULL;
272
273 r = gbo->currvalue;
274 gbo->currvalue = NULL;
275 Py_DECREF(gbo->currkey);
276 gbo->currkey = NULL;
277
278 return r;
279}
280
281static PyTypeObject _grouper_type = {
282 PyObject_HEAD_INIT(NULL)
283 0, /* ob_size */
284 "itertools._grouper", /* tp_name */
285 sizeof(_grouperobject), /* tp_basicsize */
286 0, /* tp_itemsize */
287 /* methods */
288 (destructor)_grouper_dealloc, /* tp_dealloc */
289 0, /* tp_print */
290 0, /* tp_getattr */
291 0, /* tp_setattr */
292 0, /* tp_compare */
293 0, /* tp_repr */
294 0, /* tp_as_number */
295 0, /* tp_as_sequence */
296 0, /* tp_as_mapping */
297 0, /* tp_hash */
298 0, /* tp_call */
299 0, /* tp_str */
300 PyObject_GenericGetAttr, /* tp_getattro */
301 0, /* tp_setattro */
302 0, /* tp_as_buffer */
303 Py_TPFLAGS_DEFAULT, /* tp_flags */
304 0, /* tp_doc */
305 0, /* tp_traverse */
306 0, /* tp_clear */
307 0, /* tp_richcompare */
308 0, /* tp_weaklistoffset */
309 PyObject_SelfIter, /* tp_iter */
310 (iternextfunc)_grouper_next, /* tp_iternext */
311 0, /* tp_methods */
312 0, /* tp_members */
313 0, /* tp_getset */
314 0, /* tp_base */
315 0, /* tp_dict */
316 0, /* tp_descr_get */
317 0, /* tp_descr_set */
318 0, /* tp_dictoffset */
319 0, /* tp_init */
320 0, /* tp_alloc */
321 0, /* tp_new */
322 PyObject_Del, /* tp_free */
323};
324
325
326
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
Raymond Hettingerad983e72003-11-12 14:32:26 +0000329/* The teedataobject pre-allocates space for LINKCELLS number of objects.
330 To help the object fit neatly inside cache lines (space for 16 to 32
331 pointers), the value should be a multiple of 16 minus space for
332 the other structure members including PyHEAD overhead. The larger the
333 value, the less memory overhead per object and the less time spent
334 allocating/deallocating new links. The smaller the number, the less
335 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000337#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000338
339typedef struct {
340 PyObject_HEAD
341 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342 int numread;
343 PyObject *nextlink;
344 PyObject *(values[LINKCELLS]);
345} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000346
347typedef struct {
348 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000349 teedataobject *dataobj;
350 int index;
351} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000352
Raymond Hettingerad983e72003-11-12 14:32:26 +0000353static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000354
355static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000356teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000357{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000358 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000359
Raymond Hettingerad983e72003-11-12 14:32:26 +0000360 tdo = PyObject_New(teedataobject, &teedataobject_type);
361 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000362 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000363
364 tdo->numread = 0;
365 tdo->nextlink = NULL;
366 Py_INCREF(it);
367 tdo->it = it;
368 return (PyObject *)tdo;
369}
370
371static PyObject *
372teedataobject_jumplink(teedataobject *tdo)
373{
374 if (tdo->nextlink == NULL)
375 tdo->nextlink = teedataobject_new(tdo->it);
376 Py_INCREF(tdo->nextlink);
377 return tdo->nextlink;
378}
379
380static PyObject *
381teedataobject_getitem(teedataobject *tdo, int i)
382{
383 PyObject *value;
384
385 assert(i < LINKCELLS);
386 if (i < tdo->numread)
387 value = tdo->values[i];
388 else {
389 /* this is the lead iterator, so fetch more data */
390 assert(i == tdo->numread);
391 value = PyIter_Next(tdo->it);
392 if (value == NULL)
393 return NULL;
394 tdo->numread++;
395 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000396 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000397 Py_INCREF(value);
398 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000399}
400
401static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000402teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000403{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000404 int i;
Raymond Hettinger45143692003-10-25 06:37:47 +0000405
Raymond Hettingerad983e72003-11-12 14:32:26 +0000406 for (i=0 ; i<tdo->numread ; i++)
407 Py_DECREF(tdo->values[i]);
408 Py_XDECREF(tdo->it);
409 Py_XDECREF(tdo->nextlink);
410 PyObject_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000411}
412
Raymond Hettingerad983e72003-11-12 14:32:26 +0000413PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000414
Raymond Hettingerad983e72003-11-12 14:32:26 +0000415static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000416 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000418 "itertools.tee_dataobject", /* tp_name */
419 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420 0, /* tp_itemsize */
421 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000422 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000423 0, /* tp_print */
424 0, /* tp_getattr */
425 0, /* tp_setattr */
426 0, /* tp_compare */
427 0, /* tp_repr */
428 0, /* tp_as_number */
429 0, /* tp_as_sequence */
430 0, /* tp_as_mapping */
431 0, /* tp_hash */
432 0, /* tp_call */
433 0, /* tp_str */
434 PyObject_GenericGetAttr, /* tp_getattro */
435 0, /* tp_setattro */
436 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000437 Py_TPFLAGS_DEFAULT, /* tp_flags */
438 teedataobject_doc, /* tp_doc */
439 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000440};
441
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000442
443static PyTypeObject tee_type;
444
445static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000446tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000447{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448 PyObject *value, *link;
449
450 if (to->index >= LINKCELLS) {
451 link = teedataobject_jumplink(to->dataobj);
452 Py_XDECREF(to->dataobj);
453 to->dataobj = (teedataobject *)link;
454 to->index = 0;
455 }
456 value = teedataobject_getitem(to->dataobj, to->index);
457 if (value == NULL)
458 return NULL;
459 to->index++;
460 return value;
461}
462
463static PyObject *
464tee_copy(teeobject *to)
465{
466 teeobject *newto;
467
468 newto = PyObject_New(teeobject, &tee_type);
469 if (newto == NULL)
470 return NULL;
471 Py_INCREF(to->dataobj);
472 newto->dataobj = to->dataobj;
473 newto->index = to->index;
474 return (PyObject *)newto;
475}
476
477PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
478
479static PyObject *
480tee_fromiterable(PyObject *iterable)
481{
482 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000483 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000484
485 it = PyObject_GetIter(iterable);
486 if (it == NULL)
487 return NULL;
488 if (PyObject_TypeCheck(it, &tee_type)) {
489 to = (teeobject *)tee_copy((teeobject *)it);
490 goto done;
491 }
492
493 to = PyObject_New(teeobject, &tee_type);
494 if (to == NULL)
495 goto done;
496 to->dataobj = (teedataobject *)teedataobject_new(it);
497 to->index = 0;
498done:
499 Py_XDECREF(it);
500 return (PyObject *)to;
501}
502
503static PyObject *
504tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
505{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000506 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000507
508 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
509 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000511}
512
513static void
514tee_dealloc(teeobject *to)
515{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000516 Py_XDECREF(to->dataobj);
517 PyObject_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000518}
519
Raymond Hettingerad983e72003-11-12 14:32:26 +0000520PyDoc_STRVAR(teeobject_doc,
521"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000522
Raymond Hettingerad983e72003-11-12 14:32:26 +0000523static PyMethodDef tee_methods[] = {
524 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
525 {NULL, NULL} /* sentinel */
526};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000527
528static PyTypeObject tee_type = {
529 PyObject_HEAD_INIT(NULL)
530 0, /* ob_size */
531 "itertools.tee", /* tp_name */
532 sizeof(teeobject), /* tp_basicsize */
533 0, /* tp_itemsize */
534 /* methods */
535 (destructor)tee_dealloc, /* tp_dealloc */
536 0, /* tp_print */
537 0, /* tp_getattr */
538 0, /* tp_setattr */
539 0, /* tp_compare */
540 0, /* tp_repr */
541 0, /* tp_as_number */
542 0, /* tp_as_sequence */
543 0, /* tp_as_mapping */
544 0, /* tp_hash */
545 0, /* tp_call */
546 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000547 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000548 0, /* tp_setattro */
549 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000550 Py_TPFLAGS_DEFAULT, /* tp_flags */
551 teeobject_doc, /* tp_doc */
552 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000553 0, /* tp_clear */
554 0, /* tp_richcompare */
555 0, /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000556 PyObject_SelfIter, /* tp_iter */
557 (iternextfunc)tee_next, /* tp_iternext */
558 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000559 0, /* tp_members */
560 0, /* tp_getset */
561 0, /* tp_base */
562 0, /* tp_dict */
563 0, /* tp_descr_get */
564 0, /* tp_descr_set */
565 0, /* tp_dictoffset */
566 0, /* tp_init */
567 0, /* tp_alloc */
568 tee_new, /* tp_new */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000569 PyObject_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000570};
571
Raymond Hettingerad983e72003-11-12 14:32:26 +0000572static PyObject *
573tee(PyObject *self, PyObject *args)
574{
575 int i, n=2;
576 PyObject *it, *iterable, *copyable, *result;
577
578 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
579 return NULL;
580 result = PyTuple_New(n);
581 if (result == NULL)
582 return NULL;
583 if (n == 0)
584 return result;
585 it = PyObject_GetIter(iterable);
586 if (it == NULL) {
587 Py_DECREF(result);
588 return NULL;
589 }
590 if (!PyObject_HasAttrString(it, "__copy__")) {
591 copyable = tee_fromiterable(it);
592 Py_DECREF(it);
593 if (copyable == NULL) {
594 Py_DECREF(result);
595 return NULL;
596 }
597 } else
598 copyable = it;
599 PyTuple_SET_ITEM(result, 0, copyable);
600 for (i=1 ; i<n ; i++) {
601 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
602 if (copyable == NULL) {
603 Py_DECREF(result);
604 return NULL;
605 }
606 PyTuple_SET_ITEM(result, i, copyable);
607 }
608 return result;
609}
610
611PyDoc_STRVAR(tee_doc,
612"tee(iterable, n=2) --> tuple of n independent iterators.");
613
614
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000615/* cycle object **********************************************************/
616
617typedef struct {
618 PyObject_HEAD
619 PyObject *it;
620 PyObject *saved;
621 int firstpass;
622} cycleobject;
623
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000624static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000625
626static PyObject *
627cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
628{
629 PyObject *it;
630 PyObject *iterable;
631 PyObject *saved;
632 cycleobject *lz;
633
634 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
635 return NULL;
636
637 /* Get iterator. */
638 it = PyObject_GetIter(iterable);
639 if (it == NULL)
640 return NULL;
641
642 saved = PyList_New(0);
643 if (saved == NULL) {
644 Py_DECREF(it);
645 return NULL;
646 }
647
648 /* create cycleobject structure */
649 lz = (cycleobject *)type->tp_alloc(type, 0);
650 if (lz == NULL) {
651 Py_DECREF(it);
652 Py_DECREF(saved);
653 return NULL;
654 }
655 lz->it = it;
656 lz->saved = saved;
657 lz->firstpass = 0;
658
659 return (PyObject *)lz;
660}
661
662static void
663cycle_dealloc(cycleobject *lz)
664{
665 PyObject_GC_UnTrack(lz);
666 Py_XDECREF(lz->saved);
667 Py_XDECREF(lz->it);
668 lz->ob_type->tp_free(lz);
669}
670
671static int
672cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
673{
674 int err;
675
676 if (lz->it) {
677 err = visit(lz->it, arg);
678 if (err)
679 return err;
680 }
681 if (lz->saved) {
682 err = visit(lz->saved, arg);
683 if (err)
684 return err;
685 }
686 return 0;
687}
688
689static PyObject *
690cycle_next(cycleobject *lz)
691{
692 PyObject *item;
693 PyObject *it;
694
695 while (1) {
696 item = PyIter_Next(lz->it);
697 if (item != NULL) {
698 if (!lz->firstpass)
699 PyList_Append(lz->saved, item);
700 return item;
701 }
702 if (PyList_Size(lz->saved) == 0)
703 return NULL;
704 it = PyObject_GetIter(lz->saved);
705 if (it == NULL)
706 return NULL;
707 Py_DECREF(lz->it);
708 lz->it = it;
709 lz->firstpass = 1;
710 }
711}
712
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000713PyDoc_STRVAR(cycle_doc,
714"cycle(iterable) --> cycle object\n\
715\n\
716Return elements from the iterable until it is exhausted.\n\
717Then repeat the sequence indefinitely.");
718
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000719static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000720 PyObject_HEAD_INIT(NULL)
721 0, /* ob_size */
722 "itertools.cycle", /* tp_name */
723 sizeof(cycleobject), /* tp_basicsize */
724 0, /* tp_itemsize */
725 /* methods */
726 (destructor)cycle_dealloc, /* tp_dealloc */
727 0, /* tp_print */
728 0, /* tp_getattr */
729 0, /* tp_setattr */
730 0, /* tp_compare */
731 0, /* tp_repr */
732 0, /* tp_as_number */
733 0, /* tp_as_sequence */
734 0, /* tp_as_mapping */
735 0, /* tp_hash */
736 0, /* tp_call */
737 0, /* tp_str */
738 PyObject_GenericGetAttr, /* tp_getattro */
739 0, /* tp_setattro */
740 0, /* tp_as_buffer */
741 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
742 Py_TPFLAGS_BASETYPE, /* tp_flags */
743 cycle_doc, /* tp_doc */
744 (traverseproc)cycle_traverse, /* tp_traverse */
745 0, /* tp_clear */
746 0, /* tp_richcompare */
747 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000748 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000749 (iternextfunc)cycle_next, /* tp_iternext */
750 0, /* tp_methods */
751 0, /* tp_members */
752 0, /* tp_getset */
753 0, /* tp_base */
754 0, /* tp_dict */
755 0, /* tp_descr_get */
756 0, /* tp_descr_set */
757 0, /* tp_dictoffset */
758 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000759 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760 cycle_new, /* tp_new */
761 PyObject_GC_Del, /* tp_free */
762};
763
764
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000765/* dropwhile object **********************************************************/
766
767typedef struct {
768 PyObject_HEAD
769 PyObject *func;
770 PyObject *it;
771 long start;
772} dropwhileobject;
773
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000774static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000775
776static PyObject *
777dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
778{
779 PyObject *func, *seq;
780 PyObject *it;
781 dropwhileobject *lz;
782
783 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
784 return NULL;
785
786 /* Get iterator. */
787 it = PyObject_GetIter(seq);
788 if (it == NULL)
789 return NULL;
790
791 /* create dropwhileobject structure */
792 lz = (dropwhileobject *)type->tp_alloc(type, 0);
793 if (lz == NULL) {
794 Py_DECREF(it);
795 return NULL;
796 }
797 Py_INCREF(func);
798 lz->func = func;
799 lz->it = it;
800 lz->start = 0;
801
802 return (PyObject *)lz;
803}
804
805static void
806dropwhile_dealloc(dropwhileobject *lz)
807{
808 PyObject_GC_UnTrack(lz);
809 Py_XDECREF(lz->func);
810 Py_XDECREF(lz->it);
811 lz->ob_type->tp_free(lz);
812}
813
814static int
815dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
816{
817 int err;
818
819 if (lz->it) {
820 err = visit(lz->it, arg);
821 if (err)
822 return err;
823 }
824 if (lz->func) {
825 err = visit(lz->func, arg);
826 if (err)
827 return err;
828 }
829 return 0;
830}
831
832static PyObject *
833dropwhile_next(dropwhileobject *lz)
834{
835 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000836 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000837 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000838 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000839
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000840 assert(PyIter_Check(it));
841 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000843 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000844 if (item == NULL)
845 return NULL;
846 if (lz->start == 1)
847 return item;
848
849 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
850 if (good == NULL) {
851 Py_DECREF(item);
852 return NULL;
853 }
854 ok = PyObject_IsTrue(good);
855 Py_DECREF(good);
856 if (!ok) {
857 lz->start = 1;
858 return item;
859 }
860 Py_DECREF(item);
861 }
862}
863
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864PyDoc_STRVAR(dropwhile_doc,
865"dropwhile(predicate, iterable) --> dropwhile object\n\
866\n\
867Drop items from the iterable while predicate(item) is true.\n\
868Afterwards, return every element until the iterable is exhausted.");
869
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000870static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871 PyObject_HEAD_INIT(NULL)
872 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000873 "itertools.dropwhile", /* tp_name */
874 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000875 0, /* tp_itemsize */
876 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000877 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878 0, /* tp_print */
879 0, /* tp_getattr */
880 0, /* tp_setattr */
881 0, /* tp_compare */
882 0, /* tp_repr */
883 0, /* tp_as_number */
884 0, /* tp_as_sequence */
885 0, /* tp_as_mapping */
886 0, /* tp_hash */
887 0, /* tp_call */
888 0, /* tp_str */
889 PyObject_GenericGetAttr, /* tp_getattro */
890 0, /* tp_setattro */
891 0, /* tp_as_buffer */
892 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
893 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000894 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000895 (traverseproc)dropwhile_traverse, /* tp_traverse */
896 0, /* tp_clear */
897 0, /* tp_richcompare */
898 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000899 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000900 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000901 0, /* tp_methods */
902 0, /* tp_members */
903 0, /* tp_getset */
904 0, /* tp_base */
905 0, /* tp_dict */
906 0, /* tp_descr_get */
907 0, /* tp_descr_set */
908 0, /* tp_dictoffset */
909 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000910 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000911 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912 PyObject_GC_Del, /* tp_free */
913};
914
915
916/* takewhile object **********************************************************/
917
918typedef struct {
919 PyObject_HEAD
920 PyObject *func;
921 PyObject *it;
922 long stop;
923} takewhileobject;
924
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000925static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000926
927static PyObject *
928takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
929{
930 PyObject *func, *seq;
931 PyObject *it;
932 takewhileobject *lz;
933
934 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
935 return NULL;
936
937 /* Get iterator. */
938 it = PyObject_GetIter(seq);
939 if (it == NULL)
940 return NULL;
941
942 /* create takewhileobject structure */
943 lz = (takewhileobject *)type->tp_alloc(type, 0);
944 if (lz == NULL) {
945 Py_DECREF(it);
946 return NULL;
947 }
948 Py_INCREF(func);
949 lz->func = func;
950 lz->it = it;
951 lz->stop = 0;
952
953 return (PyObject *)lz;
954}
955
956static void
957takewhile_dealloc(takewhileobject *lz)
958{
959 PyObject_GC_UnTrack(lz);
960 Py_XDECREF(lz->func);
961 Py_XDECREF(lz->it);
962 lz->ob_type->tp_free(lz);
963}
964
965static int
966takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
967{
968 int err;
969
970 if (lz->it) {
971 err = visit(lz->it, arg);
972 if (err)
973 return err;
974 }
975 if (lz->func) {
976 err = visit(lz->func, arg);
977 if (err)
978 return err;
979 }
980 return 0;
981}
982
983static PyObject *
984takewhile_next(takewhileobject *lz)
985{
986 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000987 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988 long ok;
989
990 if (lz->stop == 1)
991 return NULL;
992
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000993 assert(PyIter_Check(it));
994 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000995 if (item == NULL)
996 return NULL;
997
998 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
999 if (good == NULL) {
1000 Py_DECREF(item);
1001 return NULL;
1002 }
1003 ok = PyObject_IsTrue(good);
1004 Py_DECREF(good);
1005 if (ok)
1006 return item;
1007 Py_DECREF(item);
1008 lz->stop = 1;
1009 return NULL;
1010}
1011
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001012PyDoc_STRVAR(takewhile_doc,
1013"takewhile(predicate, iterable) --> takewhile object\n\
1014\n\
1015Return successive entries from an iterable as long as the \n\
1016predicate evaluates to true for each entry.");
1017
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001018static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001019 PyObject_HEAD_INIT(NULL)
1020 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001021 "itertools.takewhile", /* tp_name */
1022 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001023 0, /* tp_itemsize */
1024 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001025 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001026 0, /* tp_print */
1027 0, /* tp_getattr */
1028 0, /* tp_setattr */
1029 0, /* tp_compare */
1030 0, /* tp_repr */
1031 0, /* tp_as_number */
1032 0, /* tp_as_sequence */
1033 0, /* tp_as_mapping */
1034 0, /* tp_hash */
1035 0, /* tp_call */
1036 0, /* tp_str */
1037 PyObject_GenericGetAttr, /* tp_getattro */
1038 0, /* tp_setattro */
1039 0, /* tp_as_buffer */
1040 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1041 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001042 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001043 (traverseproc)takewhile_traverse, /* tp_traverse */
1044 0, /* tp_clear */
1045 0, /* tp_richcompare */
1046 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001047 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001048 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001049 0, /* tp_methods */
1050 0, /* tp_members */
1051 0, /* tp_getset */
1052 0, /* tp_base */
1053 0, /* tp_dict */
1054 0, /* tp_descr_get */
1055 0, /* tp_descr_set */
1056 0, /* tp_dictoffset */
1057 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001058 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001059 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001060 PyObject_GC_Del, /* tp_free */
1061};
1062
1063
1064/* islice object ************************************************************/
1065
1066typedef struct {
1067 PyObject_HEAD
1068 PyObject *it;
1069 long next;
1070 long stop;
1071 long step;
1072 long cnt;
1073} isliceobject;
1074
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001075static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001076
1077static PyObject *
1078islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1079{
1080 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001081 long start=0, stop=-1, step=1;
1082 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083 int numargs;
1084 isliceobject *lz;
1085
1086 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +00001087 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001088 return NULL;
1089
1090 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001091 if (a1 != Py_None) {
1092 stop = PyInt_AsLong(a1);
1093 if (stop == -1) {
1094 if (PyErr_Occurred())
1095 PyErr_Clear();
1096 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001097 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001098 return NULL;
1099 }
1100 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001101 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001102 start = PyInt_AsLong(a1);
1103 if (start == -1 && PyErr_Occurred()) {
1104 PyErr_Clear();
1105 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001106 "Start argument must be a non-negative integer.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001107 return NULL;
1108 }
1109 if (a2 != Py_None) {
1110 stop = PyInt_AsLong(a2);
1111 if (stop == -1) {
1112 if (PyErr_Occurred())
1113 PyErr_Clear();
1114 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001115 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001116 return NULL;
1117 }
1118 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119 }
1120
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001121 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001122 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001123 "Indices for islice() must be non-negative integers.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 return NULL;
1125 }
1126
1127 if (step<1) {
1128 PyErr_SetString(PyExc_ValueError,
1129 "Step must be one or larger for islice().");
1130 return NULL;
1131 }
1132
1133 /* Get iterator. */
1134 it = PyObject_GetIter(seq);
1135 if (it == NULL)
1136 return NULL;
1137
1138 /* create isliceobject structure */
1139 lz = (isliceobject *)type->tp_alloc(type, 0);
1140 if (lz == NULL) {
1141 Py_DECREF(it);
1142 return NULL;
1143 }
1144 lz->it = it;
1145 lz->next = start;
1146 lz->stop = stop;
1147 lz->step = step;
1148 lz->cnt = 0L;
1149
1150 return (PyObject *)lz;
1151}
1152
1153static void
1154islice_dealloc(isliceobject *lz)
1155{
1156 PyObject_GC_UnTrack(lz);
1157 Py_XDECREF(lz->it);
1158 lz->ob_type->tp_free(lz);
1159}
1160
1161static int
1162islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1163{
1164 if (lz->it)
1165 return visit(lz->it, arg);
1166 return 0;
1167}
1168
1169static PyObject *
1170islice_next(isliceobject *lz)
1171{
1172 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001173 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001174 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001175 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001176
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001177 assert(PyIter_Check(it));
1178 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001179 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001180 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001181 if (item == NULL)
1182 return NULL;
1183 Py_DECREF(item);
1184 lz->cnt++;
1185 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001186 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001188 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001189 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001190 if (item == NULL)
1191 return NULL;
1192 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001193 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001195 if (lz->next < oldnext) /* Check for overflow */
1196 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001197 return item;
1198}
1199
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200PyDoc_STRVAR(islice_doc,
1201"islice(iterable, [start,] stop [, step]) --> islice object\n\
1202\n\
1203Return an iterator whose next() method returns selected values from an\n\
1204iterable. If start is specified, will skip all preceding elements;\n\
1205otherwise, start defaults to zero. Step defaults to one. If\n\
1206specified as another value, step determines how many values are \n\
1207skipped between successive calls. Works like a slice() on a list\n\
1208but returns an iterator.");
1209
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001210static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211 PyObject_HEAD_INIT(NULL)
1212 0, /* ob_size */
1213 "itertools.islice", /* tp_name */
1214 sizeof(isliceobject), /* tp_basicsize */
1215 0, /* tp_itemsize */
1216 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001217 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218 0, /* tp_print */
1219 0, /* tp_getattr */
1220 0, /* tp_setattr */
1221 0, /* tp_compare */
1222 0, /* tp_repr */
1223 0, /* tp_as_number */
1224 0, /* tp_as_sequence */
1225 0, /* tp_as_mapping */
1226 0, /* tp_hash */
1227 0, /* tp_call */
1228 0, /* tp_str */
1229 PyObject_GenericGetAttr, /* tp_getattro */
1230 0, /* tp_setattro */
1231 0, /* tp_as_buffer */
1232 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1233 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001234 islice_doc, /* tp_doc */
1235 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236 0, /* tp_clear */
1237 0, /* tp_richcompare */
1238 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001239 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001240 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241 0, /* tp_methods */
1242 0, /* tp_members */
1243 0, /* tp_getset */
1244 0, /* tp_base */
1245 0, /* tp_dict */
1246 0, /* tp_descr_get */
1247 0, /* tp_descr_set */
1248 0, /* tp_dictoffset */
1249 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001250 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001251 islice_new, /* tp_new */
1252 PyObject_GC_Del, /* tp_free */
1253};
1254
1255
1256/* starmap object ************************************************************/
1257
1258typedef struct {
1259 PyObject_HEAD
1260 PyObject *func;
1261 PyObject *it;
1262} starmapobject;
1263
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001264static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001265
1266static PyObject *
1267starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1268{
1269 PyObject *func, *seq;
1270 PyObject *it;
1271 starmapobject *lz;
1272
1273 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1274 return NULL;
1275
1276 /* Get iterator. */
1277 it = PyObject_GetIter(seq);
1278 if (it == NULL)
1279 return NULL;
1280
1281 /* create starmapobject structure */
1282 lz = (starmapobject *)type->tp_alloc(type, 0);
1283 if (lz == NULL) {
1284 Py_DECREF(it);
1285 return NULL;
1286 }
1287 Py_INCREF(func);
1288 lz->func = func;
1289 lz->it = it;
1290
1291 return (PyObject *)lz;
1292}
1293
1294static void
1295starmap_dealloc(starmapobject *lz)
1296{
1297 PyObject_GC_UnTrack(lz);
1298 Py_XDECREF(lz->func);
1299 Py_XDECREF(lz->it);
1300 lz->ob_type->tp_free(lz);
1301}
1302
1303static int
1304starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1305{
1306 int err;
1307
1308 if (lz->it) {
1309 err = visit(lz->it, arg);
1310 if (err)
1311 return err;
1312 }
1313 if (lz->func) {
1314 err = visit(lz->func, arg);
1315 if (err)
1316 return err;
1317 }
1318 return 0;
1319}
1320
1321static PyObject *
1322starmap_next(starmapobject *lz)
1323{
1324 PyObject *args;
1325 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001326 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001327
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001328 assert(PyIter_Check(it));
1329 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001330 if (args == NULL)
1331 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001332 if (!PyTuple_CheckExact(args)) {
1333 Py_DECREF(args);
1334 PyErr_SetString(PyExc_TypeError,
1335 "iterator must return a tuple");
1336 return NULL;
1337 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001338 result = PyObject_Call(lz->func, args, NULL);
1339 Py_DECREF(args);
1340 return result;
1341}
1342
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001343PyDoc_STRVAR(starmap_doc,
1344"starmap(function, sequence) --> starmap object\n\
1345\n\
1346Return an iterator whose values are returned from the function evaluated\n\
1347with a argument tuple taken from the given sequence.");
1348
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001349static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350 PyObject_HEAD_INIT(NULL)
1351 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001352 "itertools.starmap", /* tp_name */
1353 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354 0, /* tp_itemsize */
1355 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001356 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001357 0, /* tp_print */
1358 0, /* tp_getattr */
1359 0, /* tp_setattr */
1360 0, /* tp_compare */
1361 0, /* tp_repr */
1362 0, /* tp_as_number */
1363 0, /* tp_as_sequence */
1364 0, /* tp_as_mapping */
1365 0, /* tp_hash */
1366 0, /* tp_call */
1367 0, /* tp_str */
1368 PyObject_GenericGetAttr, /* tp_getattro */
1369 0, /* tp_setattro */
1370 0, /* tp_as_buffer */
1371 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1372 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001373 starmap_doc, /* tp_doc */
1374 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375 0, /* tp_clear */
1376 0, /* tp_richcompare */
1377 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001378 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001379 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380 0, /* tp_methods */
1381 0, /* tp_members */
1382 0, /* tp_getset */
1383 0, /* tp_base */
1384 0, /* tp_dict */
1385 0, /* tp_descr_get */
1386 0, /* tp_descr_set */
1387 0, /* tp_dictoffset */
1388 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001389 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001390 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001391 PyObject_GC_Del, /* tp_free */
1392};
1393
1394
1395/* imap object ************************************************************/
1396
1397typedef struct {
1398 PyObject_HEAD
1399 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001400 PyObject *func;
1401} imapobject;
1402
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001403static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404
1405static PyObject *
1406imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1407{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001408 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001409 imapobject *lz;
1410 int numargs, i;
1411
1412 numargs = PyTuple_Size(args);
1413 if (numargs < 2) {
1414 PyErr_SetString(PyExc_TypeError,
1415 "imap() must have at least two arguments.");
1416 return NULL;
1417 }
1418
1419 iters = PyTuple_New(numargs-1);
1420 if (iters == NULL)
1421 return NULL;
1422
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423 for (i=1 ; i<numargs ; i++) {
1424 /* Get iterator. */
1425 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1426 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427 Py_DECREF(iters);
1428 return NULL;
1429 }
1430 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001431 }
1432
1433 /* create imapobject structure */
1434 lz = (imapobject *)type->tp_alloc(type, 0);
1435 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436 Py_DECREF(iters);
1437 return NULL;
1438 }
1439 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440 func = PyTuple_GET_ITEM(args, 0);
1441 Py_INCREF(func);
1442 lz->func = func;
1443
1444 return (PyObject *)lz;
1445}
1446
1447static void
1448imap_dealloc(imapobject *lz)
1449{
1450 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001451 Py_XDECREF(lz->iters);
1452 Py_XDECREF(lz->func);
1453 lz->ob_type->tp_free(lz);
1454}
1455
1456static int
1457imap_traverse(imapobject *lz, visitproc visit, void *arg)
1458{
1459 int err;
1460
1461 if (lz->iters) {
1462 err = visit(lz->iters, arg);
1463 if (err)
1464 return err;
1465 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466 if (lz->func) {
1467 err = visit(lz->func, arg);
1468 if (err)
1469 return err;
1470 }
1471 return 0;
1472}
1473
Raymond Hettinger2012f172003-02-07 05:32:58 +00001474/*
1475imap() is an iterator version of __builtins__.map() except that it does
1476not have the None fill-in feature. That was intentionally left out for
1477the following reasons:
1478
1479 1) Itertools are designed to be easily combined and chained together.
1480 Having all tools stop with the shortest input is a unifying principle
1481 that makes it easier to combine finite iterators (supplying data) with
1482 infinite iterators like count() and repeat() (for supplying sequential
1483 or constant arguments to a function).
1484
1485 2) In typical use cases for combining itertools, having one finite data
1486 supplier run out before another is likely to be an error condition which
1487 should not pass silently by automatically supplying None.
1488
1489 3) The use cases for automatic None fill-in are rare -- not many functions
1490 do something useful when a parameter suddenly switches type and becomes
1491 None.
1492
1493 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001494 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001495
1496 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1497*/
1498
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001499static PyObject *
1500imap_next(imapobject *lz)
1501{
1502 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001503 PyObject *argtuple;
1504 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505 int numargs, i;
1506
1507 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001508 argtuple = PyTuple_New(numargs);
1509 if (argtuple == NULL)
1510 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001511
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001512 for (i=0 ; i<numargs ; i++) {
1513 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1514 if (val == NULL) {
1515 Py_DECREF(argtuple);
1516 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001517 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001518 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001519 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001520 if (lz->func == Py_None)
1521 return argtuple;
1522 result = PyObject_Call(lz->func, argtuple, NULL);
1523 Py_DECREF(argtuple);
1524 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525}
1526
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001527PyDoc_STRVAR(imap_doc,
1528"imap(func, *iterables) --> imap object\n\
1529\n\
1530Make an iterator that computes the function using arguments from\n\
1531each of the iterables. Like map() except that it returns\n\
1532an iterator instead of a list and that it stops when the shortest\n\
1533iterable is exhausted instead of filling in None for shorter\n\
1534iterables.");
1535
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001536static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537 PyObject_HEAD_INIT(NULL)
1538 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001539 "itertools.imap", /* tp_name */
1540 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541 0, /* tp_itemsize */
1542 /* methods */
1543 (destructor)imap_dealloc, /* tp_dealloc */
1544 0, /* tp_print */
1545 0, /* tp_getattr */
1546 0, /* tp_setattr */
1547 0, /* tp_compare */
1548 0, /* tp_repr */
1549 0, /* tp_as_number */
1550 0, /* tp_as_sequence */
1551 0, /* tp_as_mapping */
1552 0, /* tp_hash */
1553 0, /* tp_call */
1554 0, /* tp_str */
1555 PyObject_GenericGetAttr, /* tp_getattro */
1556 0, /* tp_setattro */
1557 0, /* tp_as_buffer */
1558 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1559 Py_TPFLAGS_BASETYPE, /* tp_flags */
1560 imap_doc, /* tp_doc */
1561 (traverseproc)imap_traverse, /* tp_traverse */
1562 0, /* tp_clear */
1563 0, /* tp_richcompare */
1564 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001565 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001566 (iternextfunc)imap_next, /* tp_iternext */
1567 0, /* tp_methods */
1568 0, /* tp_members */
1569 0, /* tp_getset */
1570 0, /* tp_base */
1571 0, /* tp_dict */
1572 0, /* tp_descr_get */
1573 0, /* tp_descr_set */
1574 0, /* tp_dictoffset */
1575 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001576 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001577 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001578 PyObject_GC_Del, /* tp_free */
1579};
1580
1581
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001582/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583
1584typedef struct {
1585 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001586 long tuplesize;
1587 long iternum; /* which iterator is active */
1588 PyObject *ittuple; /* tuple of iterators */
1589} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001591static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001592
1593static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001594chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001596 chainobject *lz;
1597 int tuplesize = PySequence_Length(args);
1598 int i;
1599 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001600
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001601 /* obtain iterators */
1602 assert(PyTuple_Check(args));
1603 ittuple = PyTuple_New(tuplesize);
1604 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001605 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001606 for (i=0; i < tuplesize; ++i) {
1607 PyObject *item = PyTuple_GET_ITEM(args, i);
1608 PyObject *it = PyObject_GetIter(item);
1609 if (it == NULL) {
1610 if (PyErr_ExceptionMatches(PyExc_TypeError))
1611 PyErr_Format(PyExc_TypeError,
1612 "chain argument #%d must support iteration",
1613 i+1);
1614 Py_DECREF(ittuple);
1615 return NULL;
1616 }
1617 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618 }
1619
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001620 /* create chainobject structure */
1621 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001622 if (lz == NULL) {
1623 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001624 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001625 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001626
1627 lz->ittuple = ittuple;
1628 lz->iternum = 0;
1629 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001630
1631 return (PyObject *)lz;
1632}
1633
1634static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001635chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636{
1637 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001638 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639 lz->ob_type->tp_free(lz);
1640}
1641
Raymond Hettinger2012f172003-02-07 05:32:58 +00001642static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001643chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001644{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001645 if (lz->ittuple)
1646 return visit(lz->ittuple, arg);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001647 return 0;
1648}
1649
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001651chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001652{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001653 PyObject *it;
1654 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001655
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001656 while (lz->iternum < lz->tuplesize) {
1657 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1658 item = PyIter_Next(it);
1659 if (item != NULL)
1660 return item;
1661 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662 }
1663 return NULL;
1664}
1665
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001666PyDoc_STRVAR(chain_doc,
1667"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001669Return a chain object whose .next() method returns elements from the\n\
1670first iterable until it is exhausted, then elements from the next\n\
1671iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001672
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001673static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674 PyObject_HEAD_INIT(NULL)
1675 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001676 "itertools.chain", /* tp_name */
1677 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001678 0, /* tp_itemsize */
1679 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001680 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001681 0, /* tp_print */
1682 0, /* tp_getattr */
1683 0, /* tp_setattr */
1684 0, /* tp_compare */
1685 0, /* tp_repr */
1686 0, /* tp_as_number */
1687 0, /* tp_as_sequence */
1688 0, /* tp_as_mapping */
1689 0, /* tp_hash */
1690 0, /* tp_call */
1691 0, /* tp_str */
1692 PyObject_GenericGetAttr, /* tp_getattro */
1693 0, /* tp_setattro */
1694 0, /* tp_as_buffer */
1695 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1696 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001697 chain_doc, /* tp_doc */
1698 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001699 0, /* tp_clear */
1700 0, /* tp_richcompare */
1701 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001702 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001703 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001704 0, /* tp_methods */
1705 0, /* tp_members */
1706 0, /* tp_getset */
1707 0, /* tp_base */
1708 0, /* tp_dict */
1709 0, /* tp_descr_get */
1710 0, /* tp_descr_set */
1711 0, /* tp_dictoffset */
1712 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001713 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001714 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001715 PyObject_GC_Del, /* tp_free */
1716};
1717
1718
1719/* ifilter object ************************************************************/
1720
1721typedef struct {
1722 PyObject_HEAD
1723 PyObject *func;
1724 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001725} ifilterobject;
1726
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001727static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001728
1729static PyObject *
1730ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1731{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001732 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001733 PyObject *it;
1734 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735
Raymond Hettinger60eca932003-02-09 06:40:58 +00001736 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737 return NULL;
1738
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739 /* Get iterator. */
1740 it = PyObject_GetIter(seq);
1741 if (it == NULL)
1742 return NULL;
1743
1744 /* create ifilterobject structure */
1745 lz = (ifilterobject *)type->tp_alloc(type, 0);
1746 if (lz == NULL) {
1747 Py_DECREF(it);
1748 return NULL;
1749 }
1750 Py_INCREF(func);
1751 lz->func = func;
1752 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753
1754 return (PyObject *)lz;
1755}
1756
1757static void
1758ifilter_dealloc(ifilterobject *lz)
1759{
1760 PyObject_GC_UnTrack(lz);
1761 Py_XDECREF(lz->func);
1762 Py_XDECREF(lz->it);
1763 lz->ob_type->tp_free(lz);
1764}
1765
1766static int
1767ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1768{
1769 int err;
1770
1771 if (lz->it) {
1772 err = visit(lz->it, arg);
1773 if (err)
1774 return err;
1775 }
1776 if (lz->func) {
1777 err = visit(lz->func, arg);
1778 if (err)
1779 return err;
1780 }
1781 return 0;
1782}
1783
1784static PyObject *
1785ifilter_next(ifilterobject *lz)
1786{
1787 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001788 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001789 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001790 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001791
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001792 assert(PyIter_Check(it));
1793 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001794 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001795 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001796 if (item == NULL)
1797 return NULL;
1798
1799 if (lz->func == Py_None) {
1800 ok = PyObject_IsTrue(item);
1801 } else {
1802 PyObject *good;
1803 good = PyObject_CallFunctionObjArgs(lz->func,
1804 item, NULL);
1805 if (good == NULL) {
1806 Py_DECREF(item);
1807 return NULL;
1808 }
1809 ok = PyObject_IsTrue(good);
1810 Py_DECREF(good);
1811 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001812 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001813 return item;
1814 Py_DECREF(item);
1815 }
1816}
1817
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001818PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001819"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001820\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001821Return those items of sequence for which function(item) is true.\n\
1822If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001823
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001824static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001825 PyObject_HEAD_INIT(NULL)
1826 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001827 "itertools.ifilter", /* tp_name */
1828 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001829 0, /* tp_itemsize */
1830 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001831 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001832 0, /* tp_print */
1833 0, /* tp_getattr */
1834 0, /* tp_setattr */
1835 0, /* tp_compare */
1836 0, /* tp_repr */
1837 0, /* tp_as_number */
1838 0, /* tp_as_sequence */
1839 0, /* tp_as_mapping */
1840 0, /* tp_hash */
1841 0, /* tp_call */
1842 0, /* tp_str */
1843 PyObject_GenericGetAttr, /* tp_getattro */
1844 0, /* tp_setattro */
1845 0, /* tp_as_buffer */
1846 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1847 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001848 ifilter_doc, /* tp_doc */
1849 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001850 0, /* tp_clear */
1851 0, /* tp_richcompare */
1852 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001853 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001854 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001855 0, /* tp_methods */
1856 0, /* tp_members */
1857 0, /* tp_getset */
1858 0, /* tp_base */
1859 0, /* tp_dict */
1860 0, /* tp_descr_get */
1861 0, /* tp_descr_set */
1862 0, /* tp_dictoffset */
1863 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001864 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001865 ifilter_new, /* tp_new */
1866 PyObject_GC_Del, /* tp_free */
1867};
1868
1869
1870/* ifilterfalse object ************************************************************/
1871
1872typedef struct {
1873 PyObject_HEAD
1874 PyObject *func;
1875 PyObject *it;
1876} ifilterfalseobject;
1877
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001878static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001879
1880static PyObject *
1881ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1882{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001883 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001884 PyObject *it;
1885 ifilterfalseobject *lz;
1886
1887 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1888 return NULL;
1889
1890 /* Get iterator. */
1891 it = PyObject_GetIter(seq);
1892 if (it == NULL)
1893 return NULL;
1894
1895 /* create ifilterfalseobject structure */
1896 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1897 if (lz == NULL) {
1898 Py_DECREF(it);
1899 return NULL;
1900 }
1901 Py_INCREF(func);
1902 lz->func = func;
1903 lz->it = it;
1904
1905 return (PyObject *)lz;
1906}
1907
1908static void
1909ifilterfalse_dealloc(ifilterfalseobject *lz)
1910{
1911 PyObject_GC_UnTrack(lz);
1912 Py_XDECREF(lz->func);
1913 Py_XDECREF(lz->it);
1914 lz->ob_type->tp_free(lz);
1915}
1916
1917static int
1918ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1919{
1920 int err;
1921
1922 if (lz->it) {
1923 err = visit(lz->it, arg);
1924 if (err)
1925 return err;
1926 }
1927 if (lz->func) {
1928 err = visit(lz->func, arg);
1929 if (err)
1930 return err;
1931 }
1932 return 0;
1933}
1934
1935static PyObject *
1936ifilterfalse_next(ifilterfalseobject *lz)
1937{
1938 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001939 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001940 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001941 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001942
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001943 assert(PyIter_Check(it));
1944 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001945 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001946 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001947 if (item == NULL)
1948 return NULL;
1949
1950 if (lz->func == Py_None) {
1951 ok = PyObject_IsTrue(item);
1952 } else {
1953 PyObject *good;
1954 good = PyObject_CallFunctionObjArgs(lz->func,
1955 item, NULL);
1956 if (good == NULL) {
1957 Py_DECREF(item);
1958 return NULL;
1959 }
1960 ok = PyObject_IsTrue(good);
1961 Py_DECREF(good);
1962 }
1963 if (!ok)
1964 return item;
1965 Py_DECREF(item);
1966 }
1967}
1968
Raymond Hettinger60eca932003-02-09 06:40:58 +00001969PyDoc_STRVAR(ifilterfalse_doc,
1970"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1971\n\
1972Return those items of sequence for which function(item) is false.\n\
1973If function is None, return the items that are false.");
1974
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001975static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001976 PyObject_HEAD_INIT(NULL)
1977 0, /* ob_size */
1978 "itertools.ifilterfalse", /* tp_name */
1979 sizeof(ifilterfalseobject), /* tp_basicsize */
1980 0, /* tp_itemsize */
1981 /* methods */
1982 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1983 0, /* tp_print */
1984 0, /* tp_getattr */
1985 0, /* tp_setattr */
1986 0, /* tp_compare */
1987 0, /* tp_repr */
1988 0, /* tp_as_number */
1989 0, /* tp_as_sequence */
1990 0, /* tp_as_mapping */
1991 0, /* tp_hash */
1992 0, /* tp_call */
1993 0, /* tp_str */
1994 PyObject_GenericGetAttr, /* tp_getattro */
1995 0, /* tp_setattro */
1996 0, /* tp_as_buffer */
1997 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1998 Py_TPFLAGS_BASETYPE, /* tp_flags */
1999 ifilterfalse_doc, /* tp_doc */
2000 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2001 0, /* tp_clear */
2002 0, /* tp_richcompare */
2003 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002004 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002005 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2006 0, /* tp_methods */
2007 0, /* tp_members */
2008 0, /* tp_getset */
2009 0, /* tp_base */
2010 0, /* tp_dict */
2011 0, /* tp_descr_get */
2012 0, /* tp_descr_set */
2013 0, /* tp_dictoffset */
2014 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002015 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002016 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002017 PyObject_GC_Del, /* tp_free */
2018};
2019
2020
2021/* count object ************************************************************/
2022
2023typedef struct {
2024 PyObject_HEAD
2025 long cnt;
2026} countobject;
2027
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002028static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002029
2030static PyObject *
2031count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2032{
2033 countobject *lz;
2034 long cnt = 0;
2035
2036 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
2037 return NULL;
2038
2039 /* create countobject structure */
2040 lz = (countobject *)PyObject_New(countobject, &count_type);
2041 if (lz == NULL)
2042 return NULL;
2043 lz->cnt = cnt;
2044
2045 return (PyObject *)lz;
2046}
2047
2048static PyObject *
2049count_next(countobject *lz)
2050{
2051 return PyInt_FromLong(lz->cnt++);
2052}
2053
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002054static PyObject *
2055count_repr(countobject *lz)
2056{
2057 return PyString_FromFormat("count(%d)", lz->cnt);
2058}
2059
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002060PyDoc_STRVAR(count_doc,
2061"count([firstval]) --> count object\n\
2062\n\
2063Return a count object whose .next() method returns consecutive\n\
2064integers starting from zero or, if specified, from firstval.");
2065
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002066static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002067 PyObject_HEAD_INIT(NULL)
2068 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002069 "itertools.count", /* tp_name */
2070 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002071 0, /* tp_itemsize */
2072 /* methods */
2073 (destructor)PyObject_Del, /* tp_dealloc */
2074 0, /* tp_print */
2075 0, /* tp_getattr */
2076 0, /* tp_setattr */
2077 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002078 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002079 0, /* tp_as_number */
2080 0, /* tp_as_sequence */
2081 0, /* tp_as_mapping */
2082 0, /* tp_hash */
2083 0, /* tp_call */
2084 0, /* tp_str */
2085 PyObject_GenericGetAttr, /* tp_getattro */
2086 0, /* tp_setattro */
2087 0, /* tp_as_buffer */
2088 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002089 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002090 0, /* tp_traverse */
2091 0, /* tp_clear */
2092 0, /* tp_richcompare */
2093 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002094 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002095 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002096 0, /* tp_methods */
2097 0, /* tp_members */
2098 0, /* tp_getset */
2099 0, /* tp_base */
2100 0, /* tp_dict */
2101 0, /* tp_descr_get */
2102 0, /* tp_descr_set */
2103 0, /* tp_dictoffset */
2104 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002105 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002106 count_new, /* tp_new */
2107};
2108
2109
2110/* izip object ************************************************************/
2111
2112#include "Python.h"
2113
2114typedef struct {
2115 PyObject_HEAD
2116 long tuplesize;
2117 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002118 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002119} izipobject;
2120
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002121static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002122
2123static PyObject *
2124izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2125{
2126 izipobject *lz;
2127 int i;
2128 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002129 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002130 int tuplesize = PySequence_Length(args);
2131
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002132 /* args must be a tuple */
2133 assert(PyTuple_Check(args));
2134
2135 /* obtain iterators */
2136 ittuple = PyTuple_New(tuplesize);
2137 if(ittuple == NULL)
2138 return NULL;
2139 for (i=0; i < tuplesize; ++i) {
2140 PyObject *item = PyTuple_GET_ITEM(args, i);
2141 PyObject *it = PyObject_GetIter(item);
2142 if (it == NULL) {
2143 if (PyErr_ExceptionMatches(PyExc_TypeError))
2144 PyErr_Format(PyExc_TypeError,
2145 "izip argument #%d must support iteration",
2146 i+1);
2147 Py_DECREF(ittuple);
2148 return NULL;
2149 }
2150 PyTuple_SET_ITEM(ittuple, i, it);
2151 }
2152
Raymond Hettinger2012f172003-02-07 05:32:58 +00002153 /* create a result holder */
2154 result = PyTuple_New(tuplesize);
2155 if (result == NULL) {
2156 Py_DECREF(ittuple);
2157 return NULL;
2158 }
2159 for (i=0 ; i < tuplesize ; i++) {
2160 Py_INCREF(Py_None);
2161 PyTuple_SET_ITEM(result, i, Py_None);
2162 }
2163
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002164 /* create izipobject structure */
2165 lz = (izipobject *)type->tp_alloc(type, 0);
2166 if (lz == NULL) {
2167 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002168 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002169 return NULL;
2170 }
2171 lz->ittuple = ittuple;
2172 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002173 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002174
2175 return (PyObject *)lz;
2176}
2177
2178static void
2179izip_dealloc(izipobject *lz)
2180{
2181 PyObject_GC_UnTrack(lz);
2182 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002183 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002184 lz->ob_type->tp_free(lz);
2185}
2186
2187static int
2188izip_traverse(izipobject *lz, visitproc visit, void *arg)
2189{
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002190 int err;
2191
2192 if (lz->ittuple) {
2193 err = visit(lz->ittuple, arg);
2194 if (err)
2195 return err;
2196 }
2197 if (lz->result) {
2198 err = visit(lz->result, arg);
2199 if (err)
2200 return err;
2201 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002202 return 0;
2203}
2204
2205static PyObject *
2206izip_next(izipobject *lz)
2207{
2208 int i;
2209 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002210 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002211 PyObject *it;
2212 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002213 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002214
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002215 if (tuplesize == 0)
2216 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002217 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002218 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002219 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002220 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002221 assert(PyIter_Check(it));
2222 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002223 if (item == NULL) {
2224 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002225 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002226 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002227 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002228 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002229 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002230 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002231 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002232 result = PyTuple_New(tuplesize);
2233 if (result == NULL)
2234 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002235 for (i=0 ; i < tuplesize ; i++) {
2236 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002237 assert(PyIter_Check(it));
2238 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002239 if (item == NULL) {
2240 Py_DECREF(result);
2241 return NULL;
2242 }
2243 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002244 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002245 }
2246 return result;
2247}
2248
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002249PyDoc_STRVAR(izip_doc,
2250"izip(iter1 [,iter2 [...]]) --> izip object\n\
2251\n\
2252Return a izip object whose .next() method returns a tuple where\n\
2253the i-th element comes from the i-th iterable argument. The .next()\n\
2254method continues until the shortest iterable in the argument sequence\n\
2255is exhausted and then it raises StopIteration. Works like the zip()\n\
2256function but consumes less memory by returning an iterator instead of\n\
2257a list.");
2258
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002259static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002260 PyObject_HEAD_INIT(NULL)
2261 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002262 "itertools.izip", /* tp_name */
2263 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002264 0, /* tp_itemsize */
2265 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002266 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002267 0, /* tp_print */
2268 0, /* tp_getattr */
2269 0, /* tp_setattr */
2270 0, /* tp_compare */
2271 0, /* tp_repr */
2272 0, /* tp_as_number */
2273 0, /* tp_as_sequence */
2274 0, /* tp_as_mapping */
2275 0, /* tp_hash */
2276 0, /* tp_call */
2277 0, /* tp_str */
2278 PyObject_GenericGetAttr, /* tp_getattro */
2279 0, /* tp_setattro */
2280 0, /* tp_as_buffer */
2281 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2282 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002283 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002284 (traverseproc)izip_traverse, /* tp_traverse */
2285 0, /* tp_clear */
2286 0, /* tp_richcompare */
2287 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002288 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002289 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002290 0, /* tp_methods */
2291 0, /* tp_members */
2292 0, /* tp_getset */
2293 0, /* tp_base */
2294 0, /* tp_dict */
2295 0, /* tp_descr_get */
2296 0, /* tp_descr_set */
2297 0, /* tp_dictoffset */
2298 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002299 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002300 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002301 PyObject_GC_Del, /* tp_free */
2302};
2303
2304
2305/* repeat object ************************************************************/
2306
2307typedef struct {
2308 PyObject_HEAD
2309 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002310 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002311} repeatobject;
2312
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002313static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002314
2315static PyObject *
2316repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2317{
2318 repeatobject *ro;
2319 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002320 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002321
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002322 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002323 return NULL;
2324
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002325 if (PyTuple_Size(args) == 2 && cnt < 0)
2326 cnt = 0;
2327
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002328 ro = (repeatobject *)type->tp_alloc(type, 0);
2329 if (ro == NULL)
2330 return NULL;
2331 Py_INCREF(element);
2332 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002333 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002334 return (PyObject *)ro;
2335}
2336
2337static void
2338repeat_dealloc(repeatobject *ro)
2339{
2340 PyObject_GC_UnTrack(ro);
2341 Py_XDECREF(ro->element);
2342 ro->ob_type->tp_free(ro);
2343}
2344
2345static int
2346repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2347{
2348 if (ro->element)
2349 return visit(ro->element, arg);
2350 return 0;
2351}
2352
2353static PyObject *
2354repeat_next(repeatobject *ro)
2355{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002356 if (ro->cnt == 0)
2357 return NULL;
2358 if (ro->cnt > 0)
2359 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002360 Py_INCREF(ro->element);
2361 return ro->element;
2362}
2363
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002364static PyObject *
2365repeat_repr(repeatobject *ro)
2366{
2367 PyObject *result, *objrepr;
2368
2369 objrepr = PyObject_Repr(ro->element);
2370 if (objrepr == NULL)
2371 return NULL;
2372
2373 if (ro->cnt == -1)
2374 result = PyString_FromFormat("repeat(%s)",
2375 PyString_AS_STRING(objrepr));
2376 else
2377 result = PyString_FromFormat("repeat(%s, %d)",
2378 PyString_AS_STRING(objrepr), ro->cnt);
2379 Py_DECREF(objrepr);
2380 return result;
2381}
2382
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002383static int
2384repeat_len(repeatobject *ro)
2385{
2386 if (ro->cnt == -1)
2387 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2388 return (int)(ro->cnt);
2389}
2390
2391static PySequenceMethods repeat_as_sequence = {
2392 (inquiry)repeat_len, /* sq_length */
2393 0, /* sq_concat */
2394};
2395
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002396PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002397"repeat(element [,times]) -> create an iterator which returns the element\n\
2398for the specified number of times. If not specified, returns the element\n\
2399endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002400
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002401static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002402 PyObject_HEAD_INIT(NULL)
2403 0, /* ob_size */
2404 "itertools.repeat", /* tp_name */
2405 sizeof(repeatobject), /* tp_basicsize */
2406 0, /* tp_itemsize */
2407 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002408 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002409 0, /* tp_print */
2410 0, /* tp_getattr */
2411 0, /* tp_setattr */
2412 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002413 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002414 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002415 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002416 0, /* tp_as_mapping */
2417 0, /* tp_hash */
2418 0, /* tp_call */
2419 0, /* tp_str */
2420 PyObject_GenericGetAttr, /* tp_getattro */
2421 0, /* tp_setattro */
2422 0, /* tp_as_buffer */
2423 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2424 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002425 repeat_doc, /* tp_doc */
2426 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002427 0, /* tp_clear */
2428 0, /* tp_richcompare */
2429 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002430 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002431 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002432 0, /* tp_methods */
2433 0, /* tp_members */
2434 0, /* tp_getset */
2435 0, /* tp_base */
2436 0, /* tp_dict */
2437 0, /* tp_descr_get */
2438 0, /* tp_descr_set */
2439 0, /* tp_dictoffset */
2440 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002441 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002442 repeat_new, /* tp_new */
2443 PyObject_GC_Del, /* tp_free */
2444};
2445
2446
2447/* module level code ********************************************************/
2448
2449PyDoc_STRVAR(module_doc,
2450"Functional tools for creating and using iterators.\n\
2451\n\
2452Infinite iterators:\n\
2453count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002454cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002455repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002456\n\
2457Iterators terminating on the shortest input sequence:\n\
2458izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002459ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2460ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002461islice(seq, [start,] stop [, step]) --> elements from\n\
2462 seq[start:stop:step]\n\
2463imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2464starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002465tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002466chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002467takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2468dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002469groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002470");
2471
2472
Raymond Hettingerad983e72003-11-12 14:32:26 +00002473static PyMethodDef module_methods[] = {
2474 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2475 {NULL, NULL} /* sentinel */
2476};
2477
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002478PyMODINIT_FUNC
2479inititertools(void)
2480{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002481 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002482 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002483 char *name;
2484 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002485 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002486 &dropwhile_type,
2487 &takewhile_type,
2488 &islice_type,
2489 &starmap_type,
2490 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002491 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002492 &ifilter_type,
2493 &ifilterfalse_type,
2494 &count_type,
2495 &izip_type,
2496 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002497 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002498 NULL
2499 };
2500
Skip Montanarof3938fd2004-02-10 20:27:40 +00002501 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002502 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002503
Raymond Hettinger60eca932003-02-09 06:40:58 +00002504 for (i=0 ; typelist[i] != NULL ; i++) {
2505 if (PyType_Ready(typelist[i]) < 0)
2506 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002507 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002508 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002509 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002510 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002511 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002512
2513 if (PyType_Ready(&teedataobject_type) < 0)
2514 return;
2515 if (PyType_Ready(&tee_type) < 0)
2516 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002517 if (PyType_Ready(&_grouper_type) < 0)
2518 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002519}