blob: 3515bc6058bc1448ecc2809ad6bd48306f1cad9e [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 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000702 if (PyErr_Occurred()) {
703 if (PyErr_ExceptionMatches(PyExc_StopIteration))
704 PyErr_Clear();
705 else
706 return NULL;
707 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000708 if (PyList_Size(lz->saved) == 0)
709 return NULL;
710 it = PyObject_GetIter(lz->saved);
711 if (it == NULL)
712 return NULL;
713 Py_DECREF(lz->it);
714 lz->it = it;
715 lz->firstpass = 1;
716 }
717}
718
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719PyDoc_STRVAR(cycle_doc,
720"cycle(iterable) --> cycle object\n\
721\n\
722Return elements from the iterable until it is exhausted.\n\
723Then repeat the sequence indefinitely.");
724
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000725static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000726 PyObject_HEAD_INIT(NULL)
727 0, /* ob_size */
728 "itertools.cycle", /* tp_name */
729 sizeof(cycleobject), /* tp_basicsize */
730 0, /* tp_itemsize */
731 /* methods */
732 (destructor)cycle_dealloc, /* tp_dealloc */
733 0, /* tp_print */
734 0, /* tp_getattr */
735 0, /* tp_setattr */
736 0, /* tp_compare */
737 0, /* tp_repr */
738 0, /* tp_as_number */
739 0, /* tp_as_sequence */
740 0, /* tp_as_mapping */
741 0, /* tp_hash */
742 0, /* tp_call */
743 0, /* tp_str */
744 PyObject_GenericGetAttr, /* tp_getattro */
745 0, /* tp_setattro */
746 0, /* tp_as_buffer */
747 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
748 Py_TPFLAGS_BASETYPE, /* tp_flags */
749 cycle_doc, /* tp_doc */
750 (traverseproc)cycle_traverse, /* tp_traverse */
751 0, /* tp_clear */
752 0, /* tp_richcompare */
753 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000754 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000755 (iternextfunc)cycle_next, /* tp_iternext */
756 0, /* tp_methods */
757 0, /* tp_members */
758 0, /* tp_getset */
759 0, /* tp_base */
760 0, /* tp_dict */
761 0, /* tp_descr_get */
762 0, /* tp_descr_set */
763 0, /* tp_dictoffset */
764 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000765 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000766 cycle_new, /* tp_new */
767 PyObject_GC_Del, /* tp_free */
768};
769
770
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000771/* dropwhile object **********************************************************/
772
773typedef struct {
774 PyObject_HEAD
775 PyObject *func;
776 PyObject *it;
777 long start;
778} dropwhileobject;
779
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000780static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000781
782static PyObject *
783dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
784{
785 PyObject *func, *seq;
786 PyObject *it;
787 dropwhileobject *lz;
788
789 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
790 return NULL;
791
792 /* Get iterator. */
793 it = PyObject_GetIter(seq);
794 if (it == NULL)
795 return NULL;
796
797 /* create dropwhileobject structure */
798 lz = (dropwhileobject *)type->tp_alloc(type, 0);
799 if (lz == NULL) {
800 Py_DECREF(it);
801 return NULL;
802 }
803 Py_INCREF(func);
804 lz->func = func;
805 lz->it = it;
806 lz->start = 0;
807
808 return (PyObject *)lz;
809}
810
811static void
812dropwhile_dealloc(dropwhileobject *lz)
813{
814 PyObject_GC_UnTrack(lz);
815 Py_XDECREF(lz->func);
816 Py_XDECREF(lz->it);
817 lz->ob_type->tp_free(lz);
818}
819
820static int
821dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
822{
823 int err;
824
825 if (lz->it) {
826 err = visit(lz->it, arg);
827 if (err)
828 return err;
829 }
830 if (lz->func) {
831 err = visit(lz->func, arg);
832 if (err)
833 return err;
834 }
835 return 0;
836}
837
838static PyObject *
839dropwhile_next(dropwhileobject *lz)
840{
841 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000842 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000843 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000844 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000845
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000846 assert(PyIter_Check(it));
847 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000848 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000849 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000850 if (item == NULL)
851 return NULL;
852 if (lz->start == 1)
853 return item;
854
855 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
856 if (good == NULL) {
857 Py_DECREF(item);
858 return NULL;
859 }
860 ok = PyObject_IsTrue(good);
861 Py_DECREF(good);
862 if (!ok) {
863 lz->start = 1;
864 return item;
865 }
866 Py_DECREF(item);
867 }
868}
869
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000870PyDoc_STRVAR(dropwhile_doc,
871"dropwhile(predicate, iterable) --> dropwhile object\n\
872\n\
873Drop items from the iterable while predicate(item) is true.\n\
874Afterwards, return every element until the iterable is exhausted.");
875
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000876static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000877 PyObject_HEAD_INIT(NULL)
878 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000879 "itertools.dropwhile", /* tp_name */
880 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881 0, /* tp_itemsize */
882 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000883 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000884 0, /* tp_print */
885 0, /* tp_getattr */
886 0, /* tp_setattr */
887 0, /* tp_compare */
888 0, /* tp_repr */
889 0, /* tp_as_number */
890 0, /* tp_as_sequence */
891 0, /* tp_as_mapping */
892 0, /* tp_hash */
893 0, /* tp_call */
894 0, /* tp_str */
895 PyObject_GenericGetAttr, /* tp_getattro */
896 0, /* tp_setattro */
897 0, /* tp_as_buffer */
898 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
899 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000900 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000901 (traverseproc)dropwhile_traverse, /* tp_traverse */
902 0, /* tp_clear */
903 0, /* tp_richcompare */
904 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000905 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000906 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000907 0, /* tp_methods */
908 0, /* tp_members */
909 0, /* tp_getset */
910 0, /* tp_base */
911 0, /* tp_dict */
912 0, /* tp_descr_get */
913 0, /* tp_descr_set */
914 0, /* tp_dictoffset */
915 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000916 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000917 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000918 PyObject_GC_Del, /* tp_free */
919};
920
921
922/* takewhile object **********************************************************/
923
924typedef struct {
925 PyObject_HEAD
926 PyObject *func;
927 PyObject *it;
928 long stop;
929} takewhileobject;
930
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000931static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000932
933static PyObject *
934takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
935{
936 PyObject *func, *seq;
937 PyObject *it;
938 takewhileobject *lz;
939
940 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
941 return NULL;
942
943 /* Get iterator. */
944 it = PyObject_GetIter(seq);
945 if (it == NULL)
946 return NULL;
947
948 /* create takewhileobject structure */
949 lz = (takewhileobject *)type->tp_alloc(type, 0);
950 if (lz == NULL) {
951 Py_DECREF(it);
952 return NULL;
953 }
954 Py_INCREF(func);
955 lz->func = func;
956 lz->it = it;
957 lz->stop = 0;
958
959 return (PyObject *)lz;
960}
961
962static void
963takewhile_dealloc(takewhileobject *lz)
964{
965 PyObject_GC_UnTrack(lz);
966 Py_XDECREF(lz->func);
967 Py_XDECREF(lz->it);
968 lz->ob_type->tp_free(lz);
969}
970
971static int
972takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
973{
974 int err;
975
976 if (lz->it) {
977 err = visit(lz->it, arg);
978 if (err)
979 return err;
980 }
981 if (lz->func) {
982 err = visit(lz->func, arg);
983 if (err)
984 return err;
985 }
986 return 0;
987}
988
989static PyObject *
990takewhile_next(takewhileobject *lz)
991{
992 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000993 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000994 long ok;
995
996 if (lz->stop == 1)
997 return NULL;
998
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000999 assert(PyIter_Check(it));
1000 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001001 if (item == NULL)
1002 return NULL;
1003
1004 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1005 if (good == NULL) {
1006 Py_DECREF(item);
1007 return NULL;
1008 }
1009 ok = PyObject_IsTrue(good);
1010 Py_DECREF(good);
1011 if (ok)
1012 return item;
1013 Py_DECREF(item);
1014 lz->stop = 1;
1015 return NULL;
1016}
1017
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001018PyDoc_STRVAR(takewhile_doc,
1019"takewhile(predicate, iterable) --> takewhile object\n\
1020\n\
1021Return successive entries from an iterable as long as the \n\
1022predicate evaluates to true for each entry.");
1023
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001024static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025 PyObject_HEAD_INIT(NULL)
1026 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001027 "itertools.takewhile", /* tp_name */
1028 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001029 0, /* tp_itemsize */
1030 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001031 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001032 0, /* tp_print */
1033 0, /* tp_getattr */
1034 0, /* tp_setattr */
1035 0, /* tp_compare */
1036 0, /* tp_repr */
1037 0, /* tp_as_number */
1038 0, /* tp_as_sequence */
1039 0, /* tp_as_mapping */
1040 0, /* tp_hash */
1041 0, /* tp_call */
1042 0, /* tp_str */
1043 PyObject_GenericGetAttr, /* tp_getattro */
1044 0, /* tp_setattro */
1045 0, /* tp_as_buffer */
1046 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1047 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001048 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001049 (traverseproc)takewhile_traverse, /* tp_traverse */
1050 0, /* tp_clear */
1051 0, /* tp_richcompare */
1052 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001053 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001054 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001055 0, /* tp_methods */
1056 0, /* tp_members */
1057 0, /* tp_getset */
1058 0, /* tp_base */
1059 0, /* tp_dict */
1060 0, /* tp_descr_get */
1061 0, /* tp_descr_set */
1062 0, /* tp_dictoffset */
1063 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001064 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001065 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001066 PyObject_GC_Del, /* tp_free */
1067};
1068
1069
1070/* islice object ************************************************************/
1071
1072typedef struct {
1073 PyObject_HEAD
1074 PyObject *it;
1075 long next;
1076 long stop;
1077 long step;
1078 long cnt;
1079} isliceobject;
1080
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001081static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001082
1083static PyObject *
1084islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1085{
1086 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001087 long start=0, stop=-1, step=1;
1088 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 int numargs;
1090 isliceobject *lz;
1091
1092 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +00001093 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094 return NULL;
1095
1096 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001097 if (a1 != Py_None) {
1098 stop = PyInt_AsLong(a1);
1099 if (stop == -1) {
1100 if (PyErr_Occurred())
1101 PyErr_Clear();
1102 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001103 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001104 return NULL;
1105 }
1106 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001107 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001108 start = PyInt_AsLong(a1);
1109 if (start == -1 && PyErr_Occurred()) {
1110 PyErr_Clear();
1111 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001112 "Start argument must be a non-negative integer.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001113 return NULL;
1114 }
1115 if (a2 != Py_None) {
1116 stop = PyInt_AsLong(a2);
1117 if (stop == -1) {
1118 if (PyErr_Occurred())
1119 PyErr_Clear();
1120 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001121 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001122 return NULL;
1123 }
1124 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001125 }
1126
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001127 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001129 "Indices for islice() must be non-negative integers.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 return NULL;
1131 }
1132
1133 if (step<1) {
1134 PyErr_SetString(PyExc_ValueError,
1135 "Step must be one or larger for islice().");
1136 return NULL;
1137 }
1138
1139 /* Get iterator. */
1140 it = PyObject_GetIter(seq);
1141 if (it == NULL)
1142 return NULL;
1143
1144 /* create isliceobject structure */
1145 lz = (isliceobject *)type->tp_alloc(type, 0);
1146 if (lz == NULL) {
1147 Py_DECREF(it);
1148 return NULL;
1149 }
1150 lz->it = it;
1151 lz->next = start;
1152 lz->stop = stop;
1153 lz->step = step;
1154 lz->cnt = 0L;
1155
1156 return (PyObject *)lz;
1157}
1158
1159static void
1160islice_dealloc(isliceobject *lz)
1161{
1162 PyObject_GC_UnTrack(lz);
1163 Py_XDECREF(lz->it);
1164 lz->ob_type->tp_free(lz);
1165}
1166
1167static int
1168islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1169{
1170 if (lz->it)
1171 return visit(lz->it, arg);
1172 return 0;
1173}
1174
1175static PyObject *
1176islice_next(isliceobject *lz)
1177{
1178 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001179 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001180 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001181 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001183 assert(PyIter_Check(it));
1184 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001185 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001186 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187 if (item == NULL)
1188 return NULL;
1189 Py_DECREF(item);
1190 lz->cnt++;
1191 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001192 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001193 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001194 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001195 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196 if (item == NULL)
1197 return NULL;
1198 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001199 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001201 if (lz->next < oldnext) /* Check for overflow */
1202 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001203 return item;
1204}
1205
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001206PyDoc_STRVAR(islice_doc,
1207"islice(iterable, [start,] stop [, step]) --> islice object\n\
1208\n\
1209Return an iterator whose next() method returns selected values from an\n\
1210iterable. If start is specified, will skip all preceding elements;\n\
1211otherwise, start defaults to zero. Step defaults to one. If\n\
1212specified as another value, step determines how many values are \n\
1213skipped between successive calls. Works like a slice() on a list\n\
1214but returns an iterator.");
1215
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001216static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217 PyObject_HEAD_INIT(NULL)
1218 0, /* ob_size */
1219 "itertools.islice", /* tp_name */
1220 sizeof(isliceobject), /* tp_basicsize */
1221 0, /* tp_itemsize */
1222 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001223 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224 0, /* tp_print */
1225 0, /* tp_getattr */
1226 0, /* tp_setattr */
1227 0, /* tp_compare */
1228 0, /* tp_repr */
1229 0, /* tp_as_number */
1230 0, /* tp_as_sequence */
1231 0, /* tp_as_mapping */
1232 0, /* tp_hash */
1233 0, /* tp_call */
1234 0, /* tp_str */
1235 PyObject_GenericGetAttr, /* tp_getattro */
1236 0, /* tp_setattro */
1237 0, /* tp_as_buffer */
1238 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1239 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001240 islice_doc, /* tp_doc */
1241 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242 0, /* tp_clear */
1243 0, /* tp_richcompare */
1244 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001245 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001246 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001247 0, /* tp_methods */
1248 0, /* tp_members */
1249 0, /* tp_getset */
1250 0, /* tp_base */
1251 0, /* tp_dict */
1252 0, /* tp_descr_get */
1253 0, /* tp_descr_set */
1254 0, /* tp_dictoffset */
1255 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001256 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001257 islice_new, /* tp_new */
1258 PyObject_GC_Del, /* tp_free */
1259};
1260
1261
1262/* starmap object ************************************************************/
1263
1264typedef struct {
1265 PyObject_HEAD
1266 PyObject *func;
1267 PyObject *it;
1268} starmapobject;
1269
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001270static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271
1272static PyObject *
1273starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1274{
1275 PyObject *func, *seq;
1276 PyObject *it;
1277 starmapobject *lz;
1278
1279 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1280 return NULL;
1281
1282 /* Get iterator. */
1283 it = PyObject_GetIter(seq);
1284 if (it == NULL)
1285 return NULL;
1286
1287 /* create starmapobject structure */
1288 lz = (starmapobject *)type->tp_alloc(type, 0);
1289 if (lz == NULL) {
1290 Py_DECREF(it);
1291 return NULL;
1292 }
1293 Py_INCREF(func);
1294 lz->func = func;
1295 lz->it = it;
1296
1297 return (PyObject *)lz;
1298}
1299
1300static void
1301starmap_dealloc(starmapobject *lz)
1302{
1303 PyObject_GC_UnTrack(lz);
1304 Py_XDECREF(lz->func);
1305 Py_XDECREF(lz->it);
1306 lz->ob_type->tp_free(lz);
1307}
1308
1309static int
1310starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1311{
1312 int err;
1313
1314 if (lz->it) {
1315 err = visit(lz->it, arg);
1316 if (err)
1317 return err;
1318 }
1319 if (lz->func) {
1320 err = visit(lz->func, arg);
1321 if (err)
1322 return err;
1323 }
1324 return 0;
1325}
1326
1327static PyObject *
1328starmap_next(starmapobject *lz)
1329{
1330 PyObject *args;
1331 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001332 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001333
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001334 assert(PyIter_Check(it));
1335 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001336 if (args == NULL)
1337 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001338 if (!PyTuple_CheckExact(args)) {
1339 Py_DECREF(args);
1340 PyErr_SetString(PyExc_TypeError,
1341 "iterator must return a tuple");
1342 return NULL;
1343 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344 result = PyObject_Call(lz->func, args, NULL);
1345 Py_DECREF(args);
1346 return result;
1347}
1348
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001349PyDoc_STRVAR(starmap_doc,
1350"starmap(function, sequence) --> starmap object\n\
1351\n\
1352Return an iterator whose values are returned from the function evaluated\n\
1353with a argument tuple taken from the given sequence.");
1354
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001355static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001356 PyObject_HEAD_INIT(NULL)
1357 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001358 "itertools.starmap", /* tp_name */
1359 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360 0, /* tp_itemsize */
1361 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001362 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363 0, /* tp_print */
1364 0, /* tp_getattr */
1365 0, /* tp_setattr */
1366 0, /* tp_compare */
1367 0, /* tp_repr */
1368 0, /* tp_as_number */
1369 0, /* tp_as_sequence */
1370 0, /* tp_as_mapping */
1371 0, /* tp_hash */
1372 0, /* tp_call */
1373 0, /* tp_str */
1374 PyObject_GenericGetAttr, /* tp_getattro */
1375 0, /* tp_setattro */
1376 0, /* tp_as_buffer */
1377 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1378 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001379 starmap_doc, /* tp_doc */
1380 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381 0, /* tp_clear */
1382 0, /* tp_richcompare */
1383 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001384 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001385 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386 0, /* tp_methods */
1387 0, /* tp_members */
1388 0, /* tp_getset */
1389 0, /* tp_base */
1390 0, /* tp_dict */
1391 0, /* tp_descr_get */
1392 0, /* tp_descr_set */
1393 0, /* tp_dictoffset */
1394 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001395 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001396 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001397 PyObject_GC_Del, /* tp_free */
1398};
1399
1400
1401/* imap object ************************************************************/
1402
1403typedef struct {
1404 PyObject_HEAD
1405 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406 PyObject *func;
1407} imapobject;
1408
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001409static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001410
1411static PyObject *
1412imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1413{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001414 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001415 imapobject *lz;
1416 int numargs, i;
1417
1418 numargs = PyTuple_Size(args);
1419 if (numargs < 2) {
1420 PyErr_SetString(PyExc_TypeError,
1421 "imap() must have at least two arguments.");
1422 return NULL;
1423 }
1424
1425 iters = PyTuple_New(numargs-1);
1426 if (iters == NULL)
1427 return NULL;
1428
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001429 for (i=1 ; i<numargs ; i++) {
1430 /* Get iterator. */
1431 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1432 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001433 Py_DECREF(iters);
1434 return NULL;
1435 }
1436 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437 }
1438
1439 /* create imapobject structure */
1440 lz = (imapobject *)type->tp_alloc(type, 0);
1441 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001442 Py_DECREF(iters);
1443 return NULL;
1444 }
1445 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001446 func = PyTuple_GET_ITEM(args, 0);
1447 Py_INCREF(func);
1448 lz->func = func;
1449
1450 return (PyObject *)lz;
1451}
1452
1453static void
1454imap_dealloc(imapobject *lz)
1455{
1456 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457 Py_XDECREF(lz->iters);
1458 Py_XDECREF(lz->func);
1459 lz->ob_type->tp_free(lz);
1460}
1461
1462static int
1463imap_traverse(imapobject *lz, visitproc visit, void *arg)
1464{
1465 int err;
1466
1467 if (lz->iters) {
1468 err = visit(lz->iters, arg);
1469 if (err)
1470 return err;
1471 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472 if (lz->func) {
1473 err = visit(lz->func, arg);
1474 if (err)
1475 return err;
1476 }
1477 return 0;
1478}
1479
Raymond Hettinger2012f172003-02-07 05:32:58 +00001480/*
1481imap() is an iterator version of __builtins__.map() except that it does
1482not have the None fill-in feature. That was intentionally left out for
1483the following reasons:
1484
1485 1) Itertools are designed to be easily combined and chained together.
1486 Having all tools stop with the shortest input is a unifying principle
1487 that makes it easier to combine finite iterators (supplying data) with
1488 infinite iterators like count() and repeat() (for supplying sequential
1489 or constant arguments to a function).
1490
1491 2) In typical use cases for combining itertools, having one finite data
1492 supplier run out before another is likely to be an error condition which
1493 should not pass silently by automatically supplying None.
1494
1495 3) The use cases for automatic None fill-in are rare -- not many functions
1496 do something useful when a parameter suddenly switches type and becomes
1497 None.
1498
1499 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001500 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001501
1502 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1503*/
1504
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001505static PyObject *
1506imap_next(imapobject *lz)
1507{
1508 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001509 PyObject *argtuple;
1510 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001511 int numargs, i;
1512
1513 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001514 argtuple = PyTuple_New(numargs);
1515 if (argtuple == NULL)
1516 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001517
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001518 for (i=0 ; i<numargs ; i++) {
1519 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1520 if (val == NULL) {
1521 Py_DECREF(argtuple);
1522 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001524 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001526 if (lz->func == Py_None)
1527 return argtuple;
1528 result = PyObject_Call(lz->func, argtuple, NULL);
1529 Py_DECREF(argtuple);
1530 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001531}
1532
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533PyDoc_STRVAR(imap_doc,
1534"imap(func, *iterables) --> imap object\n\
1535\n\
1536Make an iterator that computes the function using arguments from\n\
1537each of the iterables. Like map() except that it returns\n\
1538an iterator instead of a list and that it stops when the shortest\n\
1539iterable is exhausted instead of filling in None for shorter\n\
1540iterables.");
1541
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001542static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543 PyObject_HEAD_INIT(NULL)
1544 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001545 "itertools.imap", /* tp_name */
1546 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001547 0, /* tp_itemsize */
1548 /* methods */
1549 (destructor)imap_dealloc, /* tp_dealloc */
1550 0, /* tp_print */
1551 0, /* tp_getattr */
1552 0, /* tp_setattr */
1553 0, /* tp_compare */
1554 0, /* tp_repr */
1555 0, /* tp_as_number */
1556 0, /* tp_as_sequence */
1557 0, /* tp_as_mapping */
1558 0, /* tp_hash */
1559 0, /* tp_call */
1560 0, /* tp_str */
1561 PyObject_GenericGetAttr, /* tp_getattro */
1562 0, /* tp_setattro */
1563 0, /* tp_as_buffer */
1564 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1565 Py_TPFLAGS_BASETYPE, /* tp_flags */
1566 imap_doc, /* tp_doc */
1567 (traverseproc)imap_traverse, /* tp_traverse */
1568 0, /* tp_clear */
1569 0, /* tp_richcompare */
1570 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001571 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001572 (iternextfunc)imap_next, /* tp_iternext */
1573 0, /* tp_methods */
1574 0, /* tp_members */
1575 0, /* tp_getset */
1576 0, /* tp_base */
1577 0, /* tp_dict */
1578 0, /* tp_descr_get */
1579 0, /* tp_descr_set */
1580 0, /* tp_dictoffset */
1581 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001582 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001583 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001584 PyObject_GC_Del, /* tp_free */
1585};
1586
1587
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001588/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589
1590typedef struct {
1591 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001592 long tuplesize;
1593 long iternum; /* which iterator is active */
1594 PyObject *ittuple; /* tuple of iterators */
1595} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001597static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001598
1599static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001600chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001601{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001602 chainobject *lz;
1603 int tuplesize = PySequence_Length(args);
1604 int i;
1605 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001607 /* obtain iterators */
1608 assert(PyTuple_Check(args));
1609 ittuple = PyTuple_New(tuplesize);
1610 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001612 for (i=0; i < tuplesize; ++i) {
1613 PyObject *item = PyTuple_GET_ITEM(args, i);
1614 PyObject *it = PyObject_GetIter(item);
1615 if (it == NULL) {
1616 if (PyErr_ExceptionMatches(PyExc_TypeError))
1617 PyErr_Format(PyExc_TypeError,
1618 "chain argument #%d must support iteration",
1619 i+1);
1620 Py_DECREF(ittuple);
1621 return NULL;
1622 }
1623 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001624 }
1625
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001626 /* create chainobject structure */
1627 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001628 if (lz == NULL) {
1629 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001630 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001631 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001632
1633 lz->ittuple = ittuple;
1634 lz->iternum = 0;
1635 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636
1637 return (PyObject *)lz;
1638}
1639
1640static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001641chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001642{
1643 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001644 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645 lz->ob_type->tp_free(lz);
1646}
1647
Raymond Hettinger2012f172003-02-07 05:32:58 +00001648static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001649chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001650{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001651 if (lz->ittuple)
1652 return visit(lz->ittuple, arg);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001653 return 0;
1654}
1655
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001656static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001657chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001659 PyObject *it;
1660 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001661
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001662 while (lz->iternum < lz->tuplesize) {
1663 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1664 item = PyIter_Next(it);
1665 if (item != NULL)
1666 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001667 if (PyErr_Occurred()) {
1668 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1669 PyErr_Clear();
1670 else
1671 return NULL;
1672 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001673 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674 }
1675 return NULL;
1676}
1677
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001678PyDoc_STRVAR(chain_doc,
1679"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001681Return a chain object whose .next() method returns elements from the\n\
1682first iterable until it is exhausted, then elements from the next\n\
1683iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001685static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001686 PyObject_HEAD_INIT(NULL)
1687 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001688 "itertools.chain", /* tp_name */
1689 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690 0, /* tp_itemsize */
1691 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001692 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001693 0, /* tp_print */
1694 0, /* tp_getattr */
1695 0, /* tp_setattr */
1696 0, /* tp_compare */
1697 0, /* tp_repr */
1698 0, /* tp_as_number */
1699 0, /* tp_as_sequence */
1700 0, /* tp_as_mapping */
1701 0, /* tp_hash */
1702 0, /* tp_call */
1703 0, /* tp_str */
1704 PyObject_GenericGetAttr, /* tp_getattro */
1705 0, /* tp_setattro */
1706 0, /* tp_as_buffer */
1707 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1708 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001709 chain_doc, /* tp_doc */
1710 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001711 0, /* tp_clear */
1712 0, /* tp_richcompare */
1713 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001714 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001715 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716 0, /* tp_methods */
1717 0, /* tp_members */
1718 0, /* tp_getset */
1719 0, /* tp_base */
1720 0, /* tp_dict */
1721 0, /* tp_descr_get */
1722 0, /* tp_descr_set */
1723 0, /* tp_dictoffset */
1724 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001725 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001726 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727 PyObject_GC_Del, /* tp_free */
1728};
1729
1730
1731/* ifilter object ************************************************************/
1732
1733typedef struct {
1734 PyObject_HEAD
1735 PyObject *func;
1736 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737} ifilterobject;
1738
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001739static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001740
1741static PyObject *
1742ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1743{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001744 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745 PyObject *it;
1746 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001747
Raymond Hettinger60eca932003-02-09 06:40:58 +00001748 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749 return NULL;
1750
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001751 /* Get iterator. */
1752 it = PyObject_GetIter(seq);
1753 if (it == NULL)
1754 return NULL;
1755
1756 /* create ifilterobject structure */
1757 lz = (ifilterobject *)type->tp_alloc(type, 0);
1758 if (lz == NULL) {
1759 Py_DECREF(it);
1760 return NULL;
1761 }
1762 Py_INCREF(func);
1763 lz->func = func;
1764 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001765
1766 return (PyObject *)lz;
1767}
1768
1769static void
1770ifilter_dealloc(ifilterobject *lz)
1771{
1772 PyObject_GC_UnTrack(lz);
1773 Py_XDECREF(lz->func);
1774 Py_XDECREF(lz->it);
1775 lz->ob_type->tp_free(lz);
1776}
1777
1778static int
1779ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1780{
1781 int err;
1782
1783 if (lz->it) {
1784 err = visit(lz->it, arg);
1785 if (err)
1786 return err;
1787 }
1788 if (lz->func) {
1789 err = visit(lz->func, arg);
1790 if (err)
1791 return err;
1792 }
1793 return 0;
1794}
1795
1796static PyObject *
1797ifilter_next(ifilterobject *lz)
1798{
1799 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001800 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001801 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001802 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001804 assert(PyIter_Check(it));
1805 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001807 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808 if (item == NULL)
1809 return NULL;
1810
1811 if (lz->func == Py_None) {
1812 ok = PyObject_IsTrue(item);
1813 } else {
1814 PyObject *good;
1815 good = PyObject_CallFunctionObjArgs(lz->func,
1816 item, NULL);
1817 if (good == NULL) {
1818 Py_DECREF(item);
1819 return NULL;
1820 }
1821 ok = PyObject_IsTrue(good);
1822 Py_DECREF(good);
1823 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001824 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001825 return item;
1826 Py_DECREF(item);
1827 }
1828}
1829
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001831"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001832\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001833Return those items of sequence for which function(item) is true.\n\
1834If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001835
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001836static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001837 PyObject_HEAD_INIT(NULL)
1838 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001839 "itertools.ifilter", /* tp_name */
1840 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001841 0, /* tp_itemsize */
1842 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001843 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001844 0, /* tp_print */
1845 0, /* tp_getattr */
1846 0, /* tp_setattr */
1847 0, /* tp_compare */
1848 0, /* tp_repr */
1849 0, /* tp_as_number */
1850 0, /* tp_as_sequence */
1851 0, /* tp_as_mapping */
1852 0, /* tp_hash */
1853 0, /* tp_call */
1854 0, /* tp_str */
1855 PyObject_GenericGetAttr, /* tp_getattro */
1856 0, /* tp_setattro */
1857 0, /* tp_as_buffer */
1858 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1859 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001860 ifilter_doc, /* tp_doc */
1861 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001862 0, /* tp_clear */
1863 0, /* tp_richcompare */
1864 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001865 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001866 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867 0, /* tp_methods */
1868 0, /* tp_members */
1869 0, /* tp_getset */
1870 0, /* tp_base */
1871 0, /* tp_dict */
1872 0, /* tp_descr_get */
1873 0, /* tp_descr_set */
1874 0, /* tp_dictoffset */
1875 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001876 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001877 ifilter_new, /* tp_new */
1878 PyObject_GC_Del, /* tp_free */
1879};
1880
1881
1882/* ifilterfalse object ************************************************************/
1883
1884typedef struct {
1885 PyObject_HEAD
1886 PyObject *func;
1887 PyObject *it;
1888} ifilterfalseobject;
1889
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001890static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001891
1892static PyObject *
1893ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1894{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001895 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001896 PyObject *it;
1897 ifilterfalseobject *lz;
1898
1899 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1900 return NULL;
1901
1902 /* Get iterator. */
1903 it = PyObject_GetIter(seq);
1904 if (it == NULL)
1905 return NULL;
1906
1907 /* create ifilterfalseobject structure */
1908 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1909 if (lz == NULL) {
1910 Py_DECREF(it);
1911 return NULL;
1912 }
1913 Py_INCREF(func);
1914 lz->func = func;
1915 lz->it = it;
1916
1917 return (PyObject *)lz;
1918}
1919
1920static void
1921ifilterfalse_dealloc(ifilterfalseobject *lz)
1922{
1923 PyObject_GC_UnTrack(lz);
1924 Py_XDECREF(lz->func);
1925 Py_XDECREF(lz->it);
1926 lz->ob_type->tp_free(lz);
1927}
1928
1929static int
1930ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1931{
1932 int err;
1933
1934 if (lz->it) {
1935 err = visit(lz->it, arg);
1936 if (err)
1937 return err;
1938 }
1939 if (lz->func) {
1940 err = visit(lz->func, arg);
1941 if (err)
1942 return err;
1943 }
1944 return 0;
1945}
1946
1947static PyObject *
1948ifilterfalse_next(ifilterfalseobject *lz)
1949{
1950 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001951 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001952 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001953 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001954
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001955 assert(PyIter_Check(it));
1956 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001957 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001958 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001959 if (item == NULL)
1960 return NULL;
1961
1962 if (lz->func == Py_None) {
1963 ok = PyObject_IsTrue(item);
1964 } else {
1965 PyObject *good;
1966 good = PyObject_CallFunctionObjArgs(lz->func,
1967 item, NULL);
1968 if (good == NULL) {
1969 Py_DECREF(item);
1970 return NULL;
1971 }
1972 ok = PyObject_IsTrue(good);
1973 Py_DECREF(good);
1974 }
1975 if (!ok)
1976 return item;
1977 Py_DECREF(item);
1978 }
1979}
1980
Raymond Hettinger60eca932003-02-09 06:40:58 +00001981PyDoc_STRVAR(ifilterfalse_doc,
1982"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1983\n\
1984Return those items of sequence for which function(item) is false.\n\
1985If function is None, return the items that are false.");
1986
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001987static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001988 PyObject_HEAD_INIT(NULL)
1989 0, /* ob_size */
1990 "itertools.ifilterfalse", /* tp_name */
1991 sizeof(ifilterfalseobject), /* tp_basicsize */
1992 0, /* tp_itemsize */
1993 /* methods */
1994 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1995 0, /* tp_print */
1996 0, /* tp_getattr */
1997 0, /* tp_setattr */
1998 0, /* tp_compare */
1999 0, /* tp_repr */
2000 0, /* tp_as_number */
2001 0, /* tp_as_sequence */
2002 0, /* tp_as_mapping */
2003 0, /* tp_hash */
2004 0, /* tp_call */
2005 0, /* tp_str */
2006 PyObject_GenericGetAttr, /* tp_getattro */
2007 0, /* tp_setattro */
2008 0, /* tp_as_buffer */
2009 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2010 Py_TPFLAGS_BASETYPE, /* tp_flags */
2011 ifilterfalse_doc, /* tp_doc */
2012 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2013 0, /* tp_clear */
2014 0, /* tp_richcompare */
2015 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002016 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002017 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2018 0, /* tp_methods */
2019 0, /* tp_members */
2020 0, /* tp_getset */
2021 0, /* tp_base */
2022 0, /* tp_dict */
2023 0, /* tp_descr_get */
2024 0, /* tp_descr_set */
2025 0, /* tp_dictoffset */
2026 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002027 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002028 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002029 PyObject_GC_Del, /* tp_free */
2030};
2031
2032
2033/* count object ************************************************************/
2034
2035typedef struct {
2036 PyObject_HEAD
2037 long cnt;
2038} countobject;
2039
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002040static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002041
2042static PyObject *
2043count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2044{
2045 countobject *lz;
2046 long cnt = 0;
2047
2048 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
2049 return NULL;
2050
2051 /* create countobject structure */
2052 lz = (countobject *)PyObject_New(countobject, &count_type);
2053 if (lz == NULL)
2054 return NULL;
2055 lz->cnt = cnt;
2056
2057 return (PyObject *)lz;
2058}
2059
2060static PyObject *
2061count_next(countobject *lz)
2062{
2063 return PyInt_FromLong(lz->cnt++);
2064}
2065
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002066static PyObject *
2067count_repr(countobject *lz)
2068{
Brett Cannon00468392004-04-13 02:43:53 +00002069 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002070}
2071
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002072PyDoc_STRVAR(count_doc,
2073"count([firstval]) --> count object\n\
2074\n\
2075Return a count object whose .next() method returns consecutive\n\
2076integers starting from zero or, if specified, from firstval.");
2077
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002078static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002079 PyObject_HEAD_INIT(NULL)
2080 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002081 "itertools.count", /* tp_name */
2082 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002083 0, /* tp_itemsize */
2084 /* methods */
2085 (destructor)PyObject_Del, /* tp_dealloc */
2086 0, /* tp_print */
2087 0, /* tp_getattr */
2088 0, /* tp_setattr */
2089 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002090 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002091 0, /* tp_as_number */
2092 0, /* tp_as_sequence */
2093 0, /* tp_as_mapping */
2094 0, /* tp_hash */
2095 0, /* tp_call */
2096 0, /* tp_str */
2097 PyObject_GenericGetAttr, /* tp_getattro */
2098 0, /* tp_setattro */
2099 0, /* tp_as_buffer */
2100 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002101 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002102 0, /* tp_traverse */
2103 0, /* tp_clear */
2104 0, /* tp_richcompare */
2105 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002106 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002107 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002108 0, /* tp_methods */
2109 0, /* tp_members */
2110 0, /* tp_getset */
2111 0, /* tp_base */
2112 0, /* tp_dict */
2113 0, /* tp_descr_get */
2114 0, /* tp_descr_set */
2115 0, /* tp_dictoffset */
2116 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002117 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002118 count_new, /* tp_new */
2119};
2120
2121
2122/* izip object ************************************************************/
2123
2124#include "Python.h"
2125
2126typedef struct {
2127 PyObject_HEAD
2128 long tuplesize;
2129 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002130 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002131} izipobject;
2132
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002133static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002134
2135static PyObject *
2136izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2137{
2138 izipobject *lz;
2139 int i;
2140 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002141 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002142 int tuplesize = PySequence_Length(args);
2143
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002144 /* args must be a tuple */
2145 assert(PyTuple_Check(args));
2146
2147 /* obtain iterators */
2148 ittuple = PyTuple_New(tuplesize);
2149 if(ittuple == NULL)
2150 return NULL;
2151 for (i=0; i < tuplesize; ++i) {
2152 PyObject *item = PyTuple_GET_ITEM(args, i);
2153 PyObject *it = PyObject_GetIter(item);
2154 if (it == NULL) {
2155 if (PyErr_ExceptionMatches(PyExc_TypeError))
2156 PyErr_Format(PyExc_TypeError,
2157 "izip argument #%d must support iteration",
2158 i+1);
2159 Py_DECREF(ittuple);
2160 return NULL;
2161 }
2162 PyTuple_SET_ITEM(ittuple, i, it);
2163 }
2164
Raymond Hettinger2012f172003-02-07 05:32:58 +00002165 /* create a result holder */
2166 result = PyTuple_New(tuplesize);
2167 if (result == NULL) {
2168 Py_DECREF(ittuple);
2169 return NULL;
2170 }
2171 for (i=0 ; i < tuplesize ; i++) {
2172 Py_INCREF(Py_None);
2173 PyTuple_SET_ITEM(result, i, Py_None);
2174 }
2175
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002176 /* create izipobject structure */
2177 lz = (izipobject *)type->tp_alloc(type, 0);
2178 if (lz == NULL) {
2179 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002180 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002181 return NULL;
2182 }
2183 lz->ittuple = ittuple;
2184 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002185 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002186
2187 return (PyObject *)lz;
2188}
2189
2190static void
2191izip_dealloc(izipobject *lz)
2192{
2193 PyObject_GC_UnTrack(lz);
2194 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002195 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002196 lz->ob_type->tp_free(lz);
2197}
2198
2199static int
2200izip_traverse(izipobject *lz, visitproc visit, void *arg)
2201{
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002202 int err;
2203
2204 if (lz->ittuple) {
2205 err = visit(lz->ittuple, arg);
2206 if (err)
2207 return err;
2208 }
2209 if (lz->result) {
2210 err = visit(lz->result, arg);
2211 if (err)
2212 return err;
2213 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002214 return 0;
2215}
2216
2217static PyObject *
2218izip_next(izipobject *lz)
2219{
2220 int i;
2221 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002222 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002223 PyObject *it;
2224 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002225 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002226
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002227 if (tuplesize == 0)
2228 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002229 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002230 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002231 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002232 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002233 assert(PyIter_Check(it));
2234 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002235 if (item == NULL) {
2236 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002237 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002238 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002239 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002240 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002241 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002242 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002243 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002244 result = PyTuple_New(tuplesize);
2245 if (result == NULL)
2246 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002247 for (i=0 ; i < tuplesize ; i++) {
2248 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002249 assert(PyIter_Check(it));
2250 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002251 if (item == NULL) {
2252 Py_DECREF(result);
2253 return NULL;
2254 }
2255 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002256 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002257 }
2258 return result;
2259}
2260
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002261PyDoc_STRVAR(izip_doc,
2262"izip(iter1 [,iter2 [...]]) --> izip object\n\
2263\n\
2264Return a izip object whose .next() method returns a tuple where\n\
2265the i-th element comes from the i-th iterable argument. The .next()\n\
2266method continues until the shortest iterable in the argument sequence\n\
2267is exhausted and then it raises StopIteration. Works like the zip()\n\
2268function but consumes less memory by returning an iterator instead of\n\
2269a list.");
2270
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002271static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002272 PyObject_HEAD_INIT(NULL)
2273 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002274 "itertools.izip", /* tp_name */
2275 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002276 0, /* tp_itemsize */
2277 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002278 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002279 0, /* tp_print */
2280 0, /* tp_getattr */
2281 0, /* tp_setattr */
2282 0, /* tp_compare */
2283 0, /* tp_repr */
2284 0, /* tp_as_number */
2285 0, /* tp_as_sequence */
2286 0, /* tp_as_mapping */
2287 0, /* tp_hash */
2288 0, /* tp_call */
2289 0, /* tp_str */
2290 PyObject_GenericGetAttr, /* tp_getattro */
2291 0, /* tp_setattro */
2292 0, /* tp_as_buffer */
2293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2294 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002295 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002296 (traverseproc)izip_traverse, /* tp_traverse */
2297 0, /* tp_clear */
2298 0, /* tp_richcompare */
2299 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002300 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002301 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002302 0, /* tp_methods */
2303 0, /* tp_members */
2304 0, /* tp_getset */
2305 0, /* tp_base */
2306 0, /* tp_dict */
2307 0, /* tp_descr_get */
2308 0, /* tp_descr_set */
2309 0, /* tp_dictoffset */
2310 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002311 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002312 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002313 PyObject_GC_Del, /* tp_free */
2314};
2315
2316
2317/* repeat object ************************************************************/
2318
2319typedef struct {
2320 PyObject_HEAD
2321 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002322 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002323} repeatobject;
2324
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002325static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002326
2327static PyObject *
2328repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2329{
2330 repeatobject *ro;
2331 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002332 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002333
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002334 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002335 return NULL;
2336
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002337 if (PyTuple_Size(args) == 2 && cnt < 0)
2338 cnt = 0;
2339
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002340 ro = (repeatobject *)type->tp_alloc(type, 0);
2341 if (ro == NULL)
2342 return NULL;
2343 Py_INCREF(element);
2344 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002345 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002346 return (PyObject *)ro;
2347}
2348
2349static void
2350repeat_dealloc(repeatobject *ro)
2351{
2352 PyObject_GC_UnTrack(ro);
2353 Py_XDECREF(ro->element);
2354 ro->ob_type->tp_free(ro);
2355}
2356
2357static int
2358repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2359{
2360 if (ro->element)
2361 return visit(ro->element, arg);
2362 return 0;
2363}
2364
2365static PyObject *
2366repeat_next(repeatobject *ro)
2367{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002368 if (ro->cnt == 0)
2369 return NULL;
2370 if (ro->cnt > 0)
2371 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002372 Py_INCREF(ro->element);
2373 return ro->element;
2374}
2375
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002376static PyObject *
2377repeat_repr(repeatobject *ro)
2378{
2379 PyObject *result, *objrepr;
2380
2381 objrepr = PyObject_Repr(ro->element);
2382 if (objrepr == NULL)
2383 return NULL;
2384
2385 if (ro->cnt == -1)
2386 result = PyString_FromFormat("repeat(%s)",
2387 PyString_AS_STRING(objrepr));
2388 else
Brett Cannon00468392004-04-13 02:43:53 +00002389 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002390 PyString_AS_STRING(objrepr), ro->cnt);
2391 Py_DECREF(objrepr);
2392 return result;
2393}
2394
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002395static int
2396repeat_len(repeatobject *ro)
2397{
2398 if (ro->cnt == -1)
2399 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2400 return (int)(ro->cnt);
2401}
2402
2403static PySequenceMethods repeat_as_sequence = {
2404 (inquiry)repeat_len, /* sq_length */
2405 0, /* sq_concat */
2406};
2407
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002408PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002409"repeat(element [,times]) -> create an iterator which returns the element\n\
2410for the specified number of times. If not specified, returns the element\n\
2411endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002412
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002413static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002414 PyObject_HEAD_INIT(NULL)
2415 0, /* ob_size */
2416 "itertools.repeat", /* tp_name */
2417 sizeof(repeatobject), /* tp_basicsize */
2418 0, /* tp_itemsize */
2419 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002420 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002421 0, /* tp_print */
2422 0, /* tp_getattr */
2423 0, /* tp_setattr */
2424 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002425 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002426 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002427 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002428 0, /* tp_as_mapping */
2429 0, /* tp_hash */
2430 0, /* tp_call */
2431 0, /* tp_str */
2432 PyObject_GenericGetAttr, /* tp_getattro */
2433 0, /* tp_setattro */
2434 0, /* tp_as_buffer */
2435 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2436 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002437 repeat_doc, /* tp_doc */
2438 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002439 0, /* tp_clear */
2440 0, /* tp_richcompare */
2441 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002442 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002443 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002444 0, /* tp_methods */
2445 0, /* tp_members */
2446 0, /* tp_getset */
2447 0, /* tp_base */
2448 0, /* tp_dict */
2449 0, /* tp_descr_get */
2450 0, /* tp_descr_set */
2451 0, /* tp_dictoffset */
2452 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002453 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002454 repeat_new, /* tp_new */
2455 PyObject_GC_Del, /* tp_free */
2456};
2457
2458
2459/* module level code ********************************************************/
2460
2461PyDoc_STRVAR(module_doc,
2462"Functional tools for creating and using iterators.\n\
2463\n\
2464Infinite iterators:\n\
2465count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002466cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002467repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002468\n\
2469Iterators terminating on the shortest input sequence:\n\
2470izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002471ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2472ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002473islice(seq, [start,] stop [, step]) --> elements from\n\
2474 seq[start:stop:step]\n\
2475imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2476starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002477tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002478chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002479takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2480dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002481groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002482");
2483
2484
Raymond Hettingerad983e72003-11-12 14:32:26 +00002485static PyMethodDef module_methods[] = {
2486 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2487 {NULL, NULL} /* sentinel */
2488};
2489
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002490PyMODINIT_FUNC
2491inititertools(void)
2492{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002493 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002494 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002495 char *name;
2496 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002497 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002498 &dropwhile_type,
2499 &takewhile_type,
2500 &islice_type,
2501 &starmap_type,
2502 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002503 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002504 &ifilter_type,
2505 &ifilterfalse_type,
2506 &count_type,
2507 &izip_type,
2508 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002509 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002510 NULL
2511 };
2512
Skip Montanarof3938fd2004-02-10 20:27:40 +00002513 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002514 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002515
Raymond Hettinger60eca932003-02-09 06:40:58 +00002516 for (i=0 ; typelist[i] != NULL ; i++) {
2517 if (PyType_Ready(typelist[i]) < 0)
2518 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002519 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002520 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002521 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002522 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002523 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002524
2525 if (PyType_Ready(&teedataobject_type) < 0)
2526 return;
2527 if (PyType_Ready(&tee_type) < 0)
2528 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002529 if (PyType_Ready(&_grouper_type) < 0)
2530 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002531}