blob: 42440df0613e5803ef0bfd317dcbb99801753d81 [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 Hettinger6a5b0272003-10-24 08:45:23 +000010/* independent iterator object supporting the tee object ***************/
11
12/* The tee object maintains a queue of data seen by the leading iterator
13 but not seen by the trailing iterator. When the leading iterator
14 gets data from PyIter_Next() it appends a copy to the inbasket stack.
15 When the trailing iterator needs data, it is popped from the outbasket
16 stack. If the outbasket stack is empty, then it is filled from the
17 inbasket (i.e. the queue is implemented using two stacks so that only
18 O(n) operations like append() and pop() are used to access data and
19 calls to reverse() never move any data element more than once).
20
21 If one of the independent iterators gets deallocated, it sets tee's
22 save_mode to zero so that future calls to PyIter_Next() stop getting
23 saved to the queue (because there is no longer a second iterator that
24 may need the data).
25*/
26
27typedef struct {
28 PyObject_HEAD
29 PyObject *it;
30 PyObject *inbasket;
31 PyObject *outbasket;
32 int save_mode;
33 int num_seen;
34} teeobject;
35
36typedef struct {
37 PyObject_HEAD
38 teeobject *tee;
39 int num_seen;
40} iiobject;
41
42static PyTypeObject ii_type;
43
44static PyObject *
45ii_next(iiobject *lz)
46{
47 teeobject *to = lz->tee;
48 PyObject *result, *tmp;
49
50 if (lz->num_seen == to->num_seen) {
51 /* This instance is leading, use iter to get more data */
52 result = PyIter_Next(to->it);
53 if (result == NULL)
54 return NULL;
55 if (to->save_mode)
56 PyList_Append(to->inbasket, result);
57 to->num_seen++;
58 lz->num_seen++;
59 return result;
60 }
61
62 /* This instance is trailing, get data from the queue */
63 if (PyList_GET_SIZE(to->outbasket) == 0) {
64 /* outbasket is empty, so refill from the inbasket */
65 tmp = to->outbasket;
66 to->outbasket = to->inbasket;
67 to->inbasket = tmp;
68 PyList_Reverse(to->outbasket);
69 assert(PyList_GET_SIZE(to->outbasket) > 0);
70 }
71
72 lz->num_seen++;
73 return PyObject_CallMethod(to->outbasket, "pop", NULL);
74}
75
76static void
77ii_dealloc(iiobject *ii)
78{
79 PyObject_GC_UnTrack(ii);
80 ii->tee->save_mode = 0; /* Stop saving data */
81 Py_XDECREF(ii->tee);
82 PyObject_GC_Del(ii);
83}
84
85static int
86ii_traverse(iiobject *ii, visitproc visit, void *arg)
87{
88 if (ii->tee)
89 return visit((PyObject *)(ii->tee), arg);
90 return 0;
91}
92
93PyDoc_STRVAR(ii_doc, "Independent iterators linked to a tee() object.");
94
95static PyTypeObject ii_type = {
96 PyObject_HEAD_INIT(&PyType_Type)
97 0, /* ob_size */
98 "itertools.independent_iterator", /* tp_name */
99 sizeof(iiobject), /* tp_basicsize */
100 0, /* tp_itemsize */
101 /* methods */
102 (destructor)ii_dealloc, /* tp_dealloc */
103 0, /* tp_print */
104 0, /* tp_getattr */
105 0, /* tp_setattr */
106 0, /* tp_compare */
107 0, /* tp_repr */
108 0, /* tp_as_number */
109 0, /* tp_as_sequence */
110 0, /* tp_as_mapping */
111 0, /* tp_hash */
112 0, /* tp_call */
113 0, /* tp_str */
114 PyObject_GenericGetAttr, /* tp_getattro */
115 0, /* tp_setattro */
116 0, /* tp_as_buffer */
117 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
118 ii_doc, /* tp_doc */
119 (traverseproc)ii_traverse, /* tp_traverse */
120 0, /* tp_clear */
121 0, /* tp_richcompare */
122 0, /* tp_weaklistoffset */
123 PyObject_SelfIter, /* tp_iter */
124 (iternextfunc)ii_next, /* tp_iternext */
125 0, /* tp_methods */
126};
127
128/* tee object **********************************************************/
129
130static PyTypeObject tee_type;
131
132static PyObject *
133tee_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
134{
135 PyObject *it = NULL;
136 PyObject *iterable;
137 PyObject *inbasket = NULL, *outbasket = NULL, *result = NULL;
138 teeobject *to = NULL;
139 int i;
140
141 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
142 return NULL;
143
144 it = PyObject_GetIter(iterable);
145 if (it == NULL) goto fail;
146
147 inbasket = PyList_New(0);
148 if (inbasket == NULL) goto fail;
149
150 outbasket = PyList_New(0);
151 if (outbasket == NULL) goto fail;
152
153 to = (teeobject *)type->tp_alloc(type, 0);
154 if (to == NULL) goto fail;
155
156 to->it = it;
157 to->inbasket = inbasket;
158 to->outbasket = outbasket;
159 to->save_mode = 1;
160 to->num_seen = 0;
161
162 /* create independent iterators */
163 result = PyTuple_New(2);
164 if (result == NULL) goto fail;
165 for (i=0 ; i<2 ; i++) {
166 iiobject *indep_it = PyObject_GC_New(iiobject, &ii_type);
167 if (indep_it == NULL) goto fail;
168 Py_INCREF(to);
169 indep_it->tee = to;
170 indep_it->num_seen = 0;
171 PyObject_GC_Track(indep_it);
172 PyTuple_SET_ITEM(result, i, (PyObject *)indep_it);
173 }
174 goto succeed;
175fail:
176 Py_XDECREF(it);
177 Py_XDECREF(inbasket);
178 Py_XDECREF(outbasket);
179 Py_XDECREF(result);
180succeed:
181 Py_XDECREF(to);
182 return result;
183}
184
185static void
186tee_dealloc(teeobject *to)
187{
188 PyObject_GC_UnTrack(to);
189 Py_XDECREF(to->inbasket);
190 Py_XDECREF(to->outbasket);
191 Py_XDECREF(to->it);
192 to->ob_type->tp_free(to);
193}
194
195static int
196tee_traverse(teeobject *to, visitproc visit, void *arg)
197{
198 int err;
199
200 if (to->it) {
201 err = visit(to->it, arg);
202 if (err)
203 return err;
204 }
205 if (to->inbasket) {
206 err = visit(to->inbasket, arg);
207 if (err)
208 return err;
209 }
210 if (to->outbasket) {
211 err = visit(to->outbasket, arg);
212 if (err)
213 return err;
214 }
215 return 0;
216}
217
218PyDoc_STRVAR(tee_doc,
219"tee(iterable) --> (it1, it2)\n\
220\n\
221Split the iterable into to independent iterables.");
222
223static PyTypeObject tee_type = {
224 PyObject_HEAD_INIT(NULL)
225 0, /* ob_size */
226 "itertools.tee", /* tp_name */
227 sizeof(teeobject), /* tp_basicsize */
228 0, /* tp_itemsize */
229 /* methods */
230 (destructor)tee_dealloc, /* tp_dealloc */
231 0, /* tp_print */
232 0, /* tp_getattr */
233 0, /* tp_setattr */
234 0, /* tp_compare */
235 0, /* tp_repr */
236 0, /* tp_as_number */
237 0, /* tp_as_sequence */
238 0, /* tp_as_mapping */
239 0, /* tp_hash */
240 0, /* tp_call */
241 0, /* tp_str */
242 PyObject_GenericGetAttr, /* tp_getattro */
243 0, /* tp_setattro */
244 0, /* tp_as_buffer */
245 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
246 Py_TPFLAGS_BASETYPE, /* tp_flags */
247 tee_doc, /* tp_doc */
248 (traverseproc)tee_traverse, /* tp_traverse */
249 0, /* tp_clear */
250 0, /* tp_richcompare */
251 0, /* tp_weaklistoffset */
252 0, /* tp_iter */
253 0, /* tp_iternext */
254 0, /* tp_methods */
255 0, /* tp_members */
256 0, /* tp_getset */
257 0, /* tp_base */
258 0, /* tp_dict */
259 0, /* tp_descr_get */
260 0, /* tp_descr_set */
261 0, /* tp_dictoffset */
262 0, /* tp_init */
263 0, /* tp_alloc */
264 tee_new, /* tp_new */
265 PyObject_GC_Del, /* tp_free */
266};
267
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000268/* cycle object **********************************************************/
269
270typedef struct {
271 PyObject_HEAD
272 PyObject *it;
273 PyObject *saved;
274 int firstpass;
275} cycleobject;
276
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000277static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000278
279static PyObject *
280cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
281{
282 PyObject *it;
283 PyObject *iterable;
284 PyObject *saved;
285 cycleobject *lz;
286
287 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
288 return NULL;
289
290 /* Get iterator. */
291 it = PyObject_GetIter(iterable);
292 if (it == NULL)
293 return NULL;
294
295 saved = PyList_New(0);
296 if (saved == NULL) {
297 Py_DECREF(it);
298 return NULL;
299 }
300
301 /* create cycleobject structure */
302 lz = (cycleobject *)type->tp_alloc(type, 0);
303 if (lz == NULL) {
304 Py_DECREF(it);
305 Py_DECREF(saved);
306 return NULL;
307 }
308 lz->it = it;
309 lz->saved = saved;
310 lz->firstpass = 0;
311
312 return (PyObject *)lz;
313}
314
315static void
316cycle_dealloc(cycleobject *lz)
317{
318 PyObject_GC_UnTrack(lz);
319 Py_XDECREF(lz->saved);
320 Py_XDECREF(lz->it);
321 lz->ob_type->tp_free(lz);
322}
323
324static int
325cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
326{
327 int err;
328
329 if (lz->it) {
330 err = visit(lz->it, arg);
331 if (err)
332 return err;
333 }
334 if (lz->saved) {
335 err = visit(lz->saved, arg);
336 if (err)
337 return err;
338 }
339 return 0;
340}
341
342static PyObject *
343cycle_next(cycleobject *lz)
344{
345 PyObject *item;
346 PyObject *it;
347
348 while (1) {
349 item = PyIter_Next(lz->it);
350 if (item != NULL) {
351 if (!lz->firstpass)
352 PyList_Append(lz->saved, item);
353 return item;
354 }
355 if (PyList_Size(lz->saved) == 0)
356 return NULL;
357 it = PyObject_GetIter(lz->saved);
358 if (it == NULL)
359 return NULL;
360 Py_DECREF(lz->it);
361 lz->it = it;
362 lz->firstpass = 1;
363 }
364}
365
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000366PyDoc_STRVAR(cycle_doc,
367"cycle(iterable) --> cycle object\n\
368\n\
369Return elements from the iterable until it is exhausted.\n\
370Then repeat the sequence indefinitely.");
371
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000372static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000373 PyObject_HEAD_INIT(NULL)
374 0, /* ob_size */
375 "itertools.cycle", /* tp_name */
376 sizeof(cycleobject), /* tp_basicsize */
377 0, /* tp_itemsize */
378 /* methods */
379 (destructor)cycle_dealloc, /* tp_dealloc */
380 0, /* tp_print */
381 0, /* tp_getattr */
382 0, /* tp_setattr */
383 0, /* tp_compare */
384 0, /* tp_repr */
385 0, /* tp_as_number */
386 0, /* tp_as_sequence */
387 0, /* tp_as_mapping */
388 0, /* tp_hash */
389 0, /* tp_call */
390 0, /* tp_str */
391 PyObject_GenericGetAttr, /* tp_getattro */
392 0, /* tp_setattro */
393 0, /* tp_as_buffer */
394 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
395 Py_TPFLAGS_BASETYPE, /* tp_flags */
396 cycle_doc, /* tp_doc */
397 (traverseproc)cycle_traverse, /* tp_traverse */
398 0, /* tp_clear */
399 0, /* tp_richcompare */
400 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000401 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000402 (iternextfunc)cycle_next, /* tp_iternext */
403 0, /* tp_methods */
404 0, /* tp_members */
405 0, /* tp_getset */
406 0, /* tp_base */
407 0, /* tp_dict */
408 0, /* tp_descr_get */
409 0, /* tp_descr_set */
410 0, /* tp_dictoffset */
411 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000412 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000413 cycle_new, /* tp_new */
414 PyObject_GC_Del, /* tp_free */
415};
416
417
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000418/* dropwhile object **********************************************************/
419
420typedef struct {
421 PyObject_HEAD
422 PyObject *func;
423 PyObject *it;
424 long start;
425} dropwhileobject;
426
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000427static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000428
429static PyObject *
430dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
431{
432 PyObject *func, *seq;
433 PyObject *it;
434 dropwhileobject *lz;
435
436 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
437 return NULL;
438
439 /* Get iterator. */
440 it = PyObject_GetIter(seq);
441 if (it == NULL)
442 return NULL;
443
444 /* create dropwhileobject structure */
445 lz = (dropwhileobject *)type->tp_alloc(type, 0);
446 if (lz == NULL) {
447 Py_DECREF(it);
448 return NULL;
449 }
450 Py_INCREF(func);
451 lz->func = func;
452 lz->it = it;
453 lz->start = 0;
454
455 return (PyObject *)lz;
456}
457
458static void
459dropwhile_dealloc(dropwhileobject *lz)
460{
461 PyObject_GC_UnTrack(lz);
462 Py_XDECREF(lz->func);
463 Py_XDECREF(lz->it);
464 lz->ob_type->tp_free(lz);
465}
466
467static int
468dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
469{
470 int err;
471
472 if (lz->it) {
473 err = visit(lz->it, arg);
474 if (err)
475 return err;
476 }
477 if (lz->func) {
478 err = visit(lz->func, arg);
479 if (err)
480 return err;
481 }
482 return 0;
483}
484
485static PyObject *
486dropwhile_next(dropwhileobject *lz)
487{
488 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000489 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000490 long ok;
491
492 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000493 assert(PyIter_Check(it));
494 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000495 if (item == NULL)
496 return NULL;
497 if (lz->start == 1)
498 return item;
499
500 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
501 if (good == NULL) {
502 Py_DECREF(item);
503 return NULL;
504 }
505 ok = PyObject_IsTrue(good);
506 Py_DECREF(good);
507 if (!ok) {
508 lz->start = 1;
509 return item;
510 }
511 Py_DECREF(item);
512 }
513}
514
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000515PyDoc_STRVAR(dropwhile_doc,
516"dropwhile(predicate, iterable) --> dropwhile object\n\
517\n\
518Drop items from the iterable while predicate(item) is true.\n\
519Afterwards, return every element until the iterable is exhausted.");
520
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000521static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000522 PyObject_HEAD_INIT(NULL)
523 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000524 "itertools.dropwhile", /* tp_name */
525 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000526 0, /* tp_itemsize */
527 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000528 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000529 0, /* tp_print */
530 0, /* tp_getattr */
531 0, /* tp_setattr */
532 0, /* tp_compare */
533 0, /* tp_repr */
534 0, /* tp_as_number */
535 0, /* tp_as_sequence */
536 0, /* tp_as_mapping */
537 0, /* tp_hash */
538 0, /* tp_call */
539 0, /* tp_str */
540 PyObject_GenericGetAttr, /* tp_getattro */
541 0, /* tp_setattro */
542 0, /* tp_as_buffer */
543 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
544 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000545 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000546 (traverseproc)dropwhile_traverse, /* tp_traverse */
547 0, /* tp_clear */
548 0, /* tp_richcompare */
549 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000550 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000551 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000552 0, /* tp_methods */
553 0, /* tp_members */
554 0, /* tp_getset */
555 0, /* tp_base */
556 0, /* tp_dict */
557 0, /* tp_descr_get */
558 0, /* tp_descr_set */
559 0, /* tp_dictoffset */
560 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000561 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000562 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000563 PyObject_GC_Del, /* tp_free */
564};
565
566
567/* takewhile object **********************************************************/
568
569typedef struct {
570 PyObject_HEAD
571 PyObject *func;
572 PyObject *it;
573 long stop;
574} takewhileobject;
575
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000576static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000577
578static PyObject *
579takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
580{
581 PyObject *func, *seq;
582 PyObject *it;
583 takewhileobject *lz;
584
585 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
586 return NULL;
587
588 /* Get iterator. */
589 it = PyObject_GetIter(seq);
590 if (it == NULL)
591 return NULL;
592
593 /* create takewhileobject structure */
594 lz = (takewhileobject *)type->tp_alloc(type, 0);
595 if (lz == NULL) {
596 Py_DECREF(it);
597 return NULL;
598 }
599 Py_INCREF(func);
600 lz->func = func;
601 lz->it = it;
602 lz->stop = 0;
603
604 return (PyObject *)lz;
605}
606
607static void
608takewhile_dealloc(takewhileobject *lz)
609{
610 PyObject_GC_UnTrack(lz);
611 Py_XDECREF(lz->func);
612 Py_XDECREF(lz->it);
613 lz->ob_type->tp_free(lz);
614}
615
616static int
617takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
618{
619 int err;
620
621 if (lz->it) {
622 err = visit(lz->it, arg);
623 if (err)
624 return err;
625 }
626 if (lz->func) {
627 err = visit(lz->func, arg);
628 if (err)
629 return err;
630 }
631 return 0;
632}
633
634static PyObject *
635takewhile_next(takewhileobject *lz)
636{
637 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000638 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000639 long ok;
640
641 if (lz->stop == 1)
642 return NULL;
643
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000644 assert(PyIter_Check(it));
645 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000646 if (item == NULL)
647 return NULL;
648
649 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
650 if (good == NULL) {
651 Py_DECREF(item);
652 return NULL;
653 }
654 ok = PyObject_IsTrue(good);
655 Py_DECREF(good);
656 if (ok)
657 return item;
658 Py_DECREF(item);
659 lz->stop = 1;
660 return NULL;
661}
662
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000663PyDoc_STRVAR(takewhile_doc,
664"takewhile(predicate, iterable) --> takewhile object\n\
665\n\
666Return successive entries from an iterable as long as the \n\
667predicate evaluates to true for each entry.");
668
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000669static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000670 PyObject_HEAD_INIT(NULL)
671 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000672 "itertools.takewhile", /* tp_name */
673 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000674 0, /* tp_itemsize */
675 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000676 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000677 0, /* tp_print */
678 0, /* tp_getattr */
679 0, /* tp_setattr */
680 0, /* tp_compare */
681 0, /* tp_repr */
682 0, /* tp_as_number */
683 0, /* tp_as_sequence */
684 0, /* tp_as_mapping */
685 0, /* tp_hash */
686 0, /* tp_call */
687 0, /* tp_str */
688 PyObject_GenericGetAttr, /* tp_getattro */
689 0, /* tp_setattro */
690 0, /* tp_as_buffer */
691 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
692 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000693 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000694 (traverseproc)takewhile_traverse, /* tp_traverse */
695 0, /* tp_clear */
696 0, /* tp_richcompare */
697 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000698 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000699 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000700 0, /* tp_methods */
701 0, /* tp_members */
702 0, /* tp_getset */
703 0, /* tp_base */
704 0, /* tp_dict */
705 0, /* tp_descr_get */
706 0, /* tp_descr_set */
707 0, /* tp_dictoffset */
708 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000709 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000710 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000711 PyObject_GC_Del, /* tp_free */
712};
713
714
715/* islice object ************************************************************/
716
717typedef struct {
718 PyObject_HEAD
719 PyObject *it;
720 long next;
721 long stop;
722 long step;
723 long cnt;
724} isliceobject;
725
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000726static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000727
728static PyObject *
729islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
730{
731 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000732 long start=0, stop=-1, step=1;
733 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000734 int numargs;
735 isliceobject *lz;
736
737 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +0000738 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000739 return NULL;
740
741 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000742 if (a1 != Py_None) {
743 stop = PyInt_AsLong(a1);
744 if (stop == -1) {
745 if (PyErr_Occurred())
746 PyErr_Clear();
747 PyErr_SetString(PyExc_ValueError,
748 "Stop argument must be an integer or None.");
749 return NULL;
750 }
751 }
Raymond Hettinger341deb72003-05-02 19:44:20 +0000752 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000753 start = PyInt_AsLong(a1);
754 if (start == -1 && PyErr_Occurred()) {
755 PyErr_Clear();
756 PyErr_SetString(PyExc_ValueError,
757 "Start argument must be an integer.");
758 return NULL;
759 }
760 if (a2 != Py_None) {
761 stop = PyInt_AsLong(a2);
762 if (stop == -1) {
763 if (PyErr_Occurred())
764 PyErr_Clear();
765 PyErr_SetString(PyExc_ValueError,
766 "Stop argument must be an integer or None.");
767 return NULL;
768 }
769 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000770 }
771
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000772 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000773 PyErr_SetString(PyExc_ValueError,
774 "Indices for islice() must be positive.");
775 return NULL;
776 }
777
778 if (step<1) {
779 PyErr_SetString(PyExc_ValueError,
780 "Step must be one or larger for islice().");
781 return NULL;
782 }
783
784 /* Get iterator. */
785 it = PyObject_GetIter(seq);
786 if (it == NULL)
787 return NULL;
788
789 /* create isliceobject structure */
790 lz = (isliceobject *)type->tp_alloc(type, 0);
791 if (lz == NULL) {
792 Py_DECREF(it);
793 return NULL;
794 }
795 lz->it = it;
796 lz->next = start;
797 lz->stop = stop;
798 lz->step = step;
799 lz->cnt = 0L;
800
801 return (PyObject *)lz;
802}
803
804static void
805islice_dealloc(isliceobject *lz)
806{
807 PyObject_GC_UnTrack(lz);
808 Py_XDECREF(lz->it);
809 lz->ob_type->tp_free(lz);
810}
811
812static int
813islice_traverse(isliceobject *lz, visitproc visit, void *arg)
814{
815 if (lz->it)
816 return visit(lz->it, arg);
817 return 0;
818}
819
820static PyObject *
821islice_next(isliceobject *lz)
822{
823 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000824 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +0000825 long oldnext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000826
827 while (lz->cnt < lz->next) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000828 assert(PyIter_Check(it));
829 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000830 if (item == NULL)
831 return NULL;
832 Py_DECREF(item);
833 lz->cnt++;
834 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000835 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000836 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000837 assert(PyIter_Check(it));
838 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000839 if (item == NULL)
840 return NULL;
841 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +0000842 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000843 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +0000844 if (lz->next < oldnext) /* Check for overflow */
845 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846 return item;
847}
848
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000849PyDoc_STRVAR(islice_doc,
850"islice(iterable, [start,] stop [, step]) --> islice object\n\
851\n\
852Return an iterator whose next() method returns selected values from an\n\
853iterable. If start is specified, will skip all preceding elements;\n\
854otherwise, start defaults to zero. Step defaults to one. If\n\
855specified as another value, step determines how many values are \n\
856skipped between successive calls. Works like a slice() on a list\n\
857but returns an iterator.");
858
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000859static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000860 PyObject_HEAD_INIT(NULL)
861 0, /* ob_size */
862 "itertools.islice", /* tp_name */
863 sizeof(isliceobject), /* tp_basicsize */
864 0, /* tp_itemsize */
865 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000866 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000867 0, /* tp_print */
868 0, /* tp_getattr */
869 0, /* tp_setattr */
870 0, /* tp_compare */
871 0, /* tp_repr */
872 0, /* tp_as_number */
873 0, /* tp_as_sequence */
874 0, /* tp_as_mapping */
875 0, /* tp_hash */
876 0, /* tp_call */
877 0, /* tp_str */
878 PyObject_GenericGetAttr, /* tp_getattro */
879 0, /* tp_setattro */
880 0, /* tp_as_buffer */
881 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
882 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000883 islice_doc, /* tp_doc */
884 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000885 0, /* tp_clear */
886 0, /* tp_richcompare */
887 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000888 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000889 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890 0, /* tp_methods */
891 0, /* tp_members */
892 0, /* tp_getset */
893 0, /* tp_base */
894 0, /* tp_dict */
895 0, /* tp_descr_get */
896 0, /* tp_descr_set */
897 0, /* tp_dictoffset */
898 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000899 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000900 islice_new, /* tp_new */
901 PyObject_GC_Del, /* tp_free */
902};
903
904
905/* starmap object ************************************************************/
906
907typedef struct {
908 PyObject_HEAD
909 PyObject *func;
910 PyObject *it;
911} starmapobject;
912
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000913static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000914
915static PyObject *
916starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
917{
918 PyObject *func, *seq;
919 PyObject *it;
920 starmapobject *lz;
921
922 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
923 return NULL;
924
925 /* Get iterator. */
926 it = PyObject_GetIter(seq);
927 if (it == NULL)
928 return NULL;
929
930 /* create starmapobject structure */
931 lz = (starmapobject *)type->tp_alloc(type, 0);
932 if (lz == NULL) {
933 Py_DECREF(it);
934 return NULL;
935 }
936 Py_INCREF(func);
937 lz->func = func;
938 lz->it = it;
939
940 return (PyObject *)lz;
941}
942
943static void
944starmap_dealloc(starmapobject *lz)
945{
946 PyObject_GC_UnTrack(lz);
947 Py_XDECREF(lz->func);
948 Py_XDECREF(lz->it);
949 lz->ob_type->tp_free(lz);
950}
951
952static int
953starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
954{
955 int err;
956
957 if (lz->it) {
958 err = visit(lz->it, arg);
959 if (err)
960 return err;
961 }
962 if (lz->func) {
963 err = visit(lz->func, arg);
964 if (err)
965 return err;
966 }
967 return 0;
968}
969
970static PyObject *
971starmap_next(starmapobject *lz)
972{
973 PyObject *args;
974 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000975 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000976
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000977 assert(PyIter_Check(it));
978 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000979 if (args == NULL)
980 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +0000981 if (!PyTuple_CheckExact(args)) {
982 Py_DECREF(args);
983 PyErr_SetString(PyExc_TypeError,
984 "iterator must return a tuple");
985 return NULL;
986 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000987 result = PyObject_Call(lz->func, args, NULL);
988 Py_DECREF(args);
989 return result;
990}
991
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000992PyDoc_STRVAR(starmap_doc,
993"starmap(function, sequence) --> starmap object\n\
994\n\
995Return an iterator whose values are returned from the function evaluated\n\
996with a argument tuple taken from the given sequence.");
997
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000998static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000999 PyObject_HEAD_INIT(NULL)
1000 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001001 "itertools.starmap", /* tp_name */
1002 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001003 0, /* tp_itemsize */
1004 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001005 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006 0, /* tp_print */
1007 0, /* tp_getattr */
1008 0, /* tp_setattr */
1009 0, /* tp_compare */
1010 0, /* tp_repr */
1011 0, /* tp_as_number */
1012 0, /* tp_as_sequence */
1013 0, /* tp_as_mapping */
1014 0, /* tp_hash */
1015 0, /* tp_call */
1016 0, /* tp_str */
1017 PyObject_GenericGetAttr, /* tp_getattro */
1018 0, /* tp_setattro */
1019 0, /* tp_as_buffer */
1020 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1021 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001022 starmap_doc, /* tp_doc */
1023 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001024 0, /* tp_clear */
1025 0, /* tp_richcompare */
1026 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001027 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001028 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001029 0, /* tp_methods */
1030 0, /* tp_members */
1031 0, /* tp_getset */
1032 0, /* tp_base */
1033 0, /* tp_dict */
1034 0, /* tp_descr_get */
1035 0, /* tp_descr_set */
1036 0, /* tp_dictoffset */
1037 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001038 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001039 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001040 PyObject_GC_Del, /* tp_free */
1041};
1042
1043
1044/* imap object ************************************************************/
1045
1046typedef struct {
1047 PyObject_HEAD
1048 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001049 PyObject *func;
1050} imapobject;
1051
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001052static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053
1054static PyObject *
1055imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1056{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001057 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001058 imapobject *lz;
1059 int numargs, i;
1060
1061 numargs = PyTuple_Size(args);
1062 if (numargs < 2) {
1063 PyErr_SetString(PyExc_TypeError,
1064 "imap() must have at least two arguments.");
1065 return NULL;
1066 }
1067
1068 iters = PyTuple_New(numargs-1);
1069 if (iters == NULL)
1070 return NULL;
1071
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072 for (i=1 ; i<numargs ; i++) {
1073 /* Get iterator. */
1074 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1075 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001076 Py_DECREF(iters);
1077 return NULL;
1078 }
1079 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001080 }
1081
1082 /* create imapobject structure */
1083 lz = (imapobject *)type->tp_alloc(type, 0);
1084 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001085 Py_DECREF(iters);
1086 return NULL;
1087 }
1088 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 func = PyTuple_GET_ITEM(args, 0);
1090 Py_INCREF(func);
1091 lz->func = func;
1092
1093 return (PyObject *)lz;
1094}
1095
1096static void
1097imap_dealloc(imapobject *lz)
1098{
1099 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100 Py_XDECREF(lz->iters);
1101 Py_XDECREF(lz->func);
1102 lz->ob_type->tp_free(lz);
1103}
1104
1105static int
1106imap_traverse(imapobject *lz, visitproc visit, void *arg)
1107{
1108 int err;
1109
1110 if (lz->iters) {
1111 err = visit(lz->iters, arg);
1112 if (err)
1113 return err;
1114 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001115 if (lz->func) {
1116 err = visit(lz->func, arg);
1117 if (err)
1118 return err;
1119 }
1120 return 0;
1121}
1122
Raymond Hettinger2012f172003-02-07 05:32:58 +00001123/*
1124imap() is an iterator version of __builtins__.map() except that it does
1125not have the None fill-in feature. That was intentionally left out for
1126the following reasons:
1127
1128 1) Itertools are designed to be easily combined and chained together.
1129 Having all tools stop with the shortest input is a unifying principle
1130 that makes it easier to combine finite iterators (supplying data) with
1131 infinite iterators like count() and repeat() (for supplying sequential
1132 or constant arguments to a function).
1133
1134 2) In typical use cases for combining itertools, having one finite data
1135 supplier run out before another is likely to be an error condition which
1136 should not pass silently by automatically supplying None.
1137
1138 3) The use cases for automatic None fill-in are rare -- not many functions
1139 do something useful when a parameter suddenly switches type and becomes
1140 None.
1141
1142 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001143 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001144
1145 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1146*/
1147
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001148static PyObject *
1149imap_next(imapobject *lz)
1150{
1151 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001152 PyObject *argtuple;
1153 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154 int numargs, i;
1155
1156 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001157 argtuple = PyTuple_New(numargs);
1158 if (argtuple == NULL)
1159 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001161 for (i=0 ; i<numargs ; i++) {
1162 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1163 if (val == NULL) {
1164 Py_DECREF(argtuple);
1165 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001166 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001167 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001168 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001169 if (lz->func == Py_None)
1170 return argtuple;
1171 result = PyObject_Call(lz->func, argtuple, NULL);
1172 Py_DECREF(argtuple);
1173 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001174}
1175
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001176PyDoc_STRVAR(imap_doc,
1177"imap(func, *iterables) --> imap object\n\
1178\n\
1179Make an iterator that computes the function using arguments from\n\
1180each of the iterables. Like map() except that it returns\n\
1181an iterator instead of a list and that it stops when the shortest\n\
1182iterable is exhausted instead of filling in None for shorter\n\
1183iterables.");
1184
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001185static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001186 PyObject_HEAD_INIT(NULL)
1187 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001188 "itertools.imap", /* tp_name */
1189 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001190 0, /* tp_itemsize */
1191 /* methods */
1192 (destructor)imap_dealloc, /* tp_dealloc */
1193 0, /* tp_print */
1194 0, /* tp_getattr */
1195 0, /* tp_setattr */
1196 0, /* tp_compare */
1197 0, /* tp_repr */
1198 0, /* tp_as_number */
1199 0, /* tp_as_sequence */
1200 0, /* tp_as_mapping */
1201 0, /* tp_hash */
1202 0, /* tp_call */
1203 0, /* tp_str */
1204 PyObject_GenericGetAttr, /* tp_getattro */
1205 0, /* tp_setattro */
1206 0, /* tp_as_buffer */
1207 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1208 Py_TPFLAGS_BASETYPE, /* tp_flags */
1209 imap_doc, /* tp_doc */
1210 (traverseproc)imap_traverse, /* tp_traverse */
1211 0, /* tp_clear */
1212 0, /* tp_richcompare */
1213 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001214 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215 (iternextfunc)imap_next, /* tp_iternext */
1216 0, /* tp_methods */
1217 0, /* tp_members */
1218 0, /* tp_getset */
1219 0, /* tp_base */
1220 0, /* tp_dict */
1221 0, /* tp_descr_get */
1222 0, /* tp_descr_set */
1223 0, /* tp_dictoffset */
1224 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001225 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001226 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001227 PyObject_GC_Del, /* tp_free */
1228};
1229
1230
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001231/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232
1233typedef struct {
1234 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001235 long tuplesize;
1236 long iternum; /* which iterator is active */
1237 PyObject *ittuple; /* tuple of iterators */
1238} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001240static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241
1242static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001243chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001245 chainobject *lz;
1246 int tuplesize = PySequence_Length(args);
1247 int i;
1248 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001249
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001250 /* obtain iterators */
1251 assert(PyTuple_Check(args));
1252 ittuple = PyTuple_New(tuplesize);
1253 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001254 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001255 for (i=0; i < tuplesize; ++i) {
1256 PyObject *item = PyTuple_GET_ITEM(args, i);
1257 PyObject *it = PyObject_GetIter(item);
1258 if (it == NULL) {
1259 if (PyErr_ExceptionMatches(PyExc_TypeError))
1260 PyErr_Format(PyExc_TypeError,
1261 "chain argument #%d must support iteration",
1262 i+1);
1263 Py_DECREF(ittuple);
1264 return NULL;
1265 }
1266 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001267 }
1268
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001269 /* create chainobject structure */
1270 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001271 if (lz == NULL) {
1272 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001273 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001274 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001275
1276 lz->ittuple = ittuple;
1277 lz->iternum = 0;
1278 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001279
1280 return (PyObject *)lz;
1281}
1282
1283static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001284chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001285{
1286 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001287 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001288 lz->ob_type->tp_free(lz);
1289}
1290
Raymond Hettinger2012f172003-02-07 05:32:58 +00001291static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001292chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001293{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001294 if (lz->ittuple)
1295 return visit(lz->ittuple, arg);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001296 return 0;
1297}
1298
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001299static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001300chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001301{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001302 PyObject *it;
1303 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001305 while (lz->iternum < lz->tuplesize) {
1306 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1307 item = PyIter_Next(it);
1308 if (item != NULL)
1309 return item;
1310 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001311 }
1312 return NULL;
1313}
1314
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001315PyDoc_STRVAR(chain_doc,
1316"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001317\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001318Return a chain object whose .next() method returns elements from the\n\
1319first iterable until it is exhausted, then elements from the next\n\
1320iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001322static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001323 PyObject_HEAD_INIT(NULL)
1324 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001325 "itertools.chain", /* tp_name */
1326 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001327 0, /* tp_itemsize */
1328 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001329 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001330 0, /* tp_print */
1331 0, /* tp_getattr */
1332 0, /* tp_setattr */
1333 0, /* tp_compare */
1334 0, /* tp_repr */
1335 0, /* tp_as_number */
1336 0, /* tp_as_sequence */
1337 0, /* tp_as_mapping */
1338 0, /* tp_hash */
1339 0, /* tp_call */
1340 0, /* tp_str */
1341 PyObject_GenericGetAttr, /* tp_getattro */
1342 0, /* tp_setattro */
1343 0, /* tp_as_buffer */
1344 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1345 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001346 chain_doc, /* tp_doc */
1347 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001348 0, /* tp_clear */
1349 0, /* tp_richcompare */
1350 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001351 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001352 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353 0, /* tp_methods */
1354 0, /* tp_members */
1355 0, /* tp_getset */
1356 0, /* tp_base */
1357 0, /* tp_dict */
1358 0, /* tp_descr_get */
1359 0, /* tp_descr_set */
1360 0, /* tp_dictoffset */
1361 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001362 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001363 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364 PyObject_GC_Del, /* tp_free */
1365};
1366
1367
1368/* ifilter object ************************************************************/
1369
1370typedef struct {
1371 PyObject_HEAD
1372 PyObject *func;
1373 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374} ifilterobject;
1375
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001376static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377
1378static PyObject *
1379ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1380{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001381 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001382 PyObject *it;
1383 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384
Raymond Hettinger60eca932003-02-09 06:40:58 +00001385 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386 return NULL;
1387
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388 /* Get iterator. */
1389 it = PyObject_GetIter(seq);
1390 if (it == NULL)
1391 return NULL;
1392
1393 /* create ifilterobject structure */
1394 lz = (ifilterobject *)type->tp_alloc(type, 0);
1395 if (lz == NULL) {
1396 Py_DECREF(it);
1397 return NULL;
1398 }
1399 Py_INCREF(func);
1400 lz->func = func;
1401 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001402
1403 return (PyObject *)lz;
1404}
1405
1406static void
1407ifilter_dealloc(ifilterobject *lz)
1408{
1409 PyObject_GC_UnTrack(lz);
1410 Py_XDECREF(lz->func);
1411 Py_XDECREF(lz->it);
1412 lz->ob_type->tp_free(lz);
1413}
1414
1415static int
1416ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1417{
1418 int err;
1419
1420 if (lz->it) {
1421 err = visit(lz->it, arg);
1422 if (err)
1423 return err;
1424 }
1425 if (lz->func) {
1426 err = visit(lz->func, arg);
1427 if (err)
1428 return err;
1429 }
1430 return 0;
1431}
1432
1433static PyObject *
1434ifilter_next(ifilterobject *lz)
1435{
1436 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001437 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438 long ok;
1439
1440 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001441 assert(PyIter_Check(it));
1442 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001443 if (item == NULL)
1444 return NULL;
1445
1446 if (lz->func == Py_None) {
1447 ok = PyObject_IsTrue(item);
1448 } else {
1449 PyObject *good;
1450 good = PyObject_CallFunctionObjArgs(lz->func,
1451 item, NULL);
1452 if (good == NULL) {
1453 Py_DECREF(item);
1454 return NULL;
1455 }
1456 ok = PyObject_IsTrue(good);
1457 Py_DECREF(good);
1458 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001459 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460 return item;
1461 Py_DECREF(item);
1462 }
1463}
1464
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001466"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001468Return those items of sequence for which function(item) is true.\n\
1469If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001471static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472 PyObject_HEAD_INIT(NULL)
1473 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001474 "itertools.ifilter", /* tp_name */
1475 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001476 0, /* tp_itemsize */
1477 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001478 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479 0, /* tp_print */
1480 0, /* tp_getattr */
1481 0, /* tp_setattr */
1482 0, /* tp_compare */
1483 0, /* tp_repr */
1484 0, /* tp_as_number */
1485 0, /* tp_as_sequence */
1486 0, /* tp_as_mapping */
1487 0, /* tp_hash */
1488 0, /* tp_call */
1489 0, /* tp_str */
1490 PyObject_GenericGetAttr, /* tp_getattro */
1491 0, /* tp_setattro */
1492 0, /* tp_as_buffer */
1493 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1494 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001495 ifilter_doc, /* tp_doc */
1496 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001497 0, /* tp_clear */
1498 0, /* tp_richcompare */
1499 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001500 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001501 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001502 0, /* tp_methods */
1503 0, /* tp_members */
1504 0, /* tp_getset */
1505 0, /* tp_base */
1506 0, /* tp_dict */
1507 0, /* tp_descr_get */
1508 0, /* tp_descr_set */
1509 0, /* tp_dictoffset */
1510 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001511 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001512 ifilter_new, /* tp_new */
1513 PyObject_GC_Del, /* tp_free */
1514};
1515
1516
1517/* ifilterfalse object ************************************************************/
1518
1519typedef struct {
1520 PyObject_HEAD
1521 PyObject *func;
1522 PyObject *it;
1523} ifilterfalseobject;
1524
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001525static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001526
1527static PyObject *
1528ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1529{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001530 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001531 PyObject *it;
1532 ifilterfalseobject *lz;
1533
1534 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1535 return NULL;
1536
1537 /* Get iterator. */
1538 it = PyObject_GetIter(seq);
1539 if (it == NULL)
1540 return NULL;
1541
1542 /* create ifilterfalseobject structure */
1543 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1544 if (lz == NULL) {
1545 Py_DECREF(it);
1546 return NULL;
1547 }
1548 Py_INCREF(func);
1549 lz->func = func;
1550 lz->it = it;
1551
1552 return (PyObject *)lz;
1553}
1554
1555static void
1556ifilterfalse_dealloc(ifilterfalseobject *lz)
1557{
1558 PyObject_GC_UnTrack(lz);
1559 Py_XDECREF(lz->func);
1560 Py_XDECREF(lz->it);
1561 lz->ob_type->tp_free(lz);
1562}
1563
1564static int
1565ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1566{
1567 int err;
1568
1569 if (lz->it) {
1570 err = visit(lz->it, arg);
1571 if (err)
1572 return err;
1573 }
1574 if (lz->func) {
1575 err = visit(lz->func, arg);
1576 if (err)
1577 return err;
1578 }
1579 return 0;
1580}
1581
1582static PyObject *
1583ifilterfalse_next(ifilterfalseobject *lz)
1584{
1585 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001586 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001587 long ok;
1588
1589 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001590 assert(PyIter_Check(it));
1591 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001592 if (item == NULL)
1593 return NULL;
1594
1595 if (lz->func == Py_None) {
1596 ok = PyObject_IsTrue(item);
1597 } else {
1598 PyObject *good;
1599 good = PyObject_CallFunctionObjArgs(lz->func,
1600 item, NULL);
1601 if (good == NULL) {
1602 Py_DECREF(item);
1603 return NULL;
1604 }
1605 ok = PyObject_IsTrue(good);
1606 Py_DECREF(good);
1607 }
1608 if (!ok)
1609 return item;
1610 Py_DECREF(item);
1611 }
1612}
1613
Raymond Hettinger60eca932003-02-09 06:40:58 +00001614PyDoc_STRVAR(ifilterfalse_doc,
1615"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1616\n\
1617Return those items of sequence for which function(item) is false.\n\
1618If function is None, return the items that are false.");
1619
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001620static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001621 PyObject_HEAD_INIT(NULL)
1622 0, /* ob_size */
1623 "itertools.ifilterfalse", /* tp_name */
1624 sizeof(ifilterfalseobject), /* tp_basicsize */
1625 0, /* tp_itemsize */
1626 /* methods */
1627 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1628 0, /* tp_print */
1629 0, /* tp_getattr */
1630 0, /* tp_setattr */
1631 0, /* tp_compare */
1632 0, /* tp_repr */
1633 0, /* tp_as_number */
1634 0, /* tp_as_sequence */
1635 0, /* tp_as_mapping */
1636 0, /* tp_hash */
1637 0, /* tp_call */
1638 0, /* tp_str */
1639 PyObject_GenericGetAttr, /* tp_getattro */
1640 0, /* tp_setattro */
1641 0, /* tp_as_buffer */
1642 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1643 Py_TPFLAGS_BASETYPE, /* tp_flags */
1644 ifilterfalse_doc, /* tp_doc */
1645 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1646 0, /* tp_clear */
1647 0, /* tp_richcompare */
1648 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001649 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001650 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1651 0, /* tp_methods */
1652 0, /* tp_members */
1653 0, /* tp_getset */
1654 0, /* tp_base */
1655 0, /* tp_dict */
1656 0, /* tp_descr_get */
1657 0, /* tp_descr_set */
1658 0, /* tp_dictoffset */
1659 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001660 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001661 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662 PyObject_GC_Del, /* tp_free */
1663};
1664
1665
1666/* count object ************************************************************/
1667
1668typedef struct {
1669 PyObject_HEAD
1670 long cnt;
1671} countobject;
1672
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001673static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674
1675static PyObject *
1676count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1677{
1678 countobject *lz;
1679 long cnt = 0;
1680
1681 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1682 return NULL;
1683
1684 /* create countobject structure */
1685 lz = (countobject *)PyObject_New(countobject, &count_type);
1686 if (lz == NULL)
1687 return NULL;
1688 lz->cnt = cnt;
1689
1690 return (PyObject *)lz;
1691}
1692
1693static PyObject *
1694count_next(countobject *lz)
1695{
1696 return PyInt_FromLong(lz->cnt++);
1697}
1698
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001699PyDoc_STRVAR(count_doc,
1700"count([firstval]) --> count object\n\
1701\n\
1702Return a count object whose .next() method returns consecutive\n\
1703integers starting from zero or, if specified, from firstval.");
1704
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001705static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001706 PyObject_HEAD_INIT(NULL)
1707 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001708 "itertools.count", /* tp_name */
1709 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001710 0, /* tp_itemsize */
1711 /* methods */
1712 (destructor)PyObject_Del, /* tp_dealloc */
1713 0, /* tp_print */
1714 0, /* tp_getattr */
1715 0, /* tp_setattr */
1716 0, /* tp_compare */
1717 0, /* tp_repr */
1718 0, /* tp_as_number */
1719 0, /* tp_as_sequence */
1720 0, /* tp_as_mapping */
1721 0, /* tp_hash */
1722 0, /* tp_call */
1723 0, /* tp_str */
1724 PyObject_GenericGetAttr, /* tp_getattro */
1725 0, /* tp_setattro */
1726 0, /* tp_as_buffer */
1727 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001728 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001729 0, /* tp_traverse */
1730 0, /* tp_clear */
1731 0, /* tp_richcompare */
1732 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001733 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001734 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735 0, /* tp_methods */
1736 0, /* tp_members */
1737 0, /* tp_getset */
1738 0, /* tp_base */
1739 0, /* tp_dict */
1740 0, /* tp_descr_get */
1741 0, /* tp_descr_set */
1742 0, /* tp_dictoffset */
1743 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001744 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745 count_new, /* tp_new */
1746};
1747
1748
1749/* izip object ************************************************************/
1750
1751#include "Python.h"
1752
1753typedef struct {
1754 PyObject_HEAD
1755 long tuplesize;
1756 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00001757 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758} izipobject;
1759
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001760static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001761
1762static PyObject *
1763izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1764{
1765 izipobject *lz;
1766 int i;
1767 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00001768 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001769 int tuplesize = PySequence_Length(args);
1770
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001771 /* args must be a tuple */
1772 assert(PyTuple_Check(args));
1773
1774 /* obtain iterators */
1775 ittuple = PyTuple_New(tuplesize);
1776 if(ittuple == NULL)
1777 return NULL;
1778 for (i=0; i < tuplesize; ++i) {
1779 PyObject *item = PyTuple_GET_ITEM(args, i);
1780 PyObject *it = PyObject_GetIter(item);
1781 if (it == NULL) {
1782 if (PyErr_ExceptionMatches(PyExc_TypeError))
1783 PyErr_Format(PyExc_TypeError,
1784 "izip argument #%d must support iteration",
1785 i+1);
1786 Py_DECREF(ittuple);
1787 return NULL;
1788 }
1789 PyTuple_SET_ITEM(ittuple, i, it);
1790 }
1791
Raymond Hettinger2012f172003-02-07 05:32:58 +00001792 /* create a result holder */
1793 result = PyTuple_New(tuplesize);
1794 if (result == NULL) {
1795 Py_DECREF(ittuple);
1796 return NULL;
1797 }
1798 for (i=0 ; i < tuplesize ; i++) {
1799 Py_INCREF(Py_None);
1800 PyTuple_SET_ITEM(result, i, Py_None);
1801 }
1802
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001803 /* create izipobject structure */
1804 lz = (izipobject *)type->tp_alloc(type, 0);
1805 if (lz == NULL) {
1806 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001807 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808 return NULL;
1809 }
1810 lz->ittuple = ittuple;
1811 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001812 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001813
1814 return (PyObject *)lz;
1815}
1816
1817static void
1818izip_dealloc(izipobject *lz)
1819{
1820 PyObject_GC_UnTrack(lz);
1821 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001822 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001823 lz->ob_type->tp_free(lz);
1824}
1825
1826static int
1827izip_traverse(izipobject *lz, visitproc visit, void *arg)
1828{
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001829 int err;
1830
1831 if (lz->ittuple) {
1832 err = visit(lz->ittuple, arg);
1833 if (err)
1834 return err;
1835 }
1836 if (lz->result) {
1837 err = visit(lz->result, arg);
1838 if (err)
1839 return err;
1840 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001841 return 0;
1842}
1843
1844static PyObject *
1845izip_next(izipobject *lz)
1846{
1847 int i;
1848 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001849 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001850 PyObject *it;
1851 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00001852 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001853
Raymond Hettingerb5a42082003-08-08 05:10:41 +00001854 if (tuplesize == 0)
1855 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001856 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001857 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001858 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001859 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001860 assert(PyIter_Check(it));
1861 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001862 if (item == NULL) {
1863 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001864 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00001865 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00001866 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001867 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00001868 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001869 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00001870 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00001871 result = PyTuple_New(tuplesize);
1872 if (result == NULL)
1873 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001874 for (i=0 ; i < tuplesize ; i++) {
1875 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001876 assert(PyIter_Check(it));
1877 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001878 if (item == NULL) {
1879 Py_DECREF(result);
1880 return NULL;
1881 }
1882 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001883 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001884 }
1885 return result;
1886}
1887
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001888PyDoc_STRVAR(izip_doc,
1889"izip(iter1 [,iter2 [...]]) --> izip object\n\
1890\n\
1891Return a izip object whose .next() method returns a tuple where\n\
1892the i-th element comes from the i-th iterable argument. The .next()\n\
1893method continues until the shortest iterable in the argument sequence\n\
1894is exhausted and then it raises StopIteration. Works like the zip()\n\
1895function but consumes less memory by returning an iterator instead of\n\
1896a list.");
1897
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001898static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001899 PyObject_HEAD_INIT(NULL)
1900 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001901 "itertools.izip", /* tp_name */
1902 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001903 0, /* tp_itemsize */
1904 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001905 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001906 0, /* tp_print */
1907 0, /* tp_getattr */
1908 0, /* tp_setattr */
1909 0, /* tp_compare */
1910 0, /* tp_repr */
1911 0, /* tp_as_number */
1912 0, /* tp_as_sequence */
1913 0, /* tp_as_mapping */
1914 0, /* tp_hash */
1915 0, /* tp_call */
1916 0, /* tp_str */
1917 PyObject_GenericGetAttr, /* tp_getattro */
1918 0, /* tp_setattro */
1919 0, /* tp_as_buffer */
1920 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1921 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001922 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001923 (traverseproc)izip_traverse, /* tp_traverse */
1924 0, /* tp_clear */
1925 0, /* tp_richcompare */
1926 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001927 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001928 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001929 0, /* tp_methods */
1930 0, /* tp_members */
1931 0, /* tp_getset */
1932 0, /* tp_base */
1933 0, /* tp_dict */
1934 0, /* tp_descr_get */
1935 0, /* tp_descr_set */
1936 0, /* tp_dictoffset */
1937 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001938 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001939 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001940 PyObject_GC_Del, /* tp_free */
1941};
1942
1943
1944/* repeat object ************************************************************/
1945
1946typedef struct {
1947 PyObject_HEAD
1948 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001949 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001950} repeatobject;
1951
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001952static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001953
1954static PyObject *
1955repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1956{
1957 repeatobject *ro;
1958 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001959 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001960
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001961 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001962 return NULL;
1963
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001964 if (PyTuple_Size(args) == 2 && cnt < 0)
1965 cnt = 0;
1966
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001967 ro = (repeatobject *)type->tp_alloc(type, 0);
1968 if (ro == NULL)
1969 return NULL;
1970 Py_INCREF(element);
1971 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001972 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001973 return (PyObject *)ro;
1974}
1975
1976static void
1977repeat_dealloc(repeatobject *ro)
1978{
1979 PyObject_GC_UnTrack(ro);
1980 Py_XDECREF(ro->element);
1981 ro->ob_type->tp_free(ro);
1982}
1983
1984static int
1985repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
1986{
1987 if (ro->element)
1988 return visit(ro->element, arg);
1989 return 0;
1990}
1991
1992static PyObject *
1993repeat_next(repeatobject *ro)
1994{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001995 if (ro->cnt == 0)
1996 return NULL;
1997 if (ro->cnt > 0)
1998 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001999 Py_INCREF(ro->element);
2000 return ro->element;
2001}
2002
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002003PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002004"repeat(element [,times]) -> create an iterator which returns the element\n\
2005for the specified number of times. If not specified, returns the element\n\
2006endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002007
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002008static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002009 PyObject_HEAD_INIT(NULL)
2010 0, /* ob_size */
2011 "itertools.repeat", /* tp_name */
2012 sizeof(repeatobject), /* tp_basicsize */
2013 0, /* tp_itemsize */
2014 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002015 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002016 0, /* tp_print */
2017 0, /* tp_getattr */
2018 0, /* tp_setattr */
2019 0, /* tp_compare */
2020 0, /* tp_repr */
2021 0, /* tp_as_number */
2022 0, /* tp_as_sequence */
2023 0, /* tp_as_mapping */
2024 0, /* tp_hash */
2025 0, /* tp_call */
2026 0, /* tp_str */
2027 PyObject_GenericGetAttr, /* tp_getattro */
2028 0, /* tp_setattro */
2029 0, /* tp_as_buffer */
2030 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2031 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002032 repeat_doc, /* tp_doc */
2033 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002034 0, /* tp_clear */
2035 0, /* tp_richcompare */
2036 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002037 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002038 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002039 0, /* tp_methods */
2040 0, /* tp_members */
2041 0, /* tp_getset */
2042 0, /* tp_base */
2043 0, /* tp_dict */
2044 0, /* tp_descr_get */
2045 0, /* tp_descr_set */
2046 0, /* tp_dictoffset */
2047 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002048 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002049 repeat_new, /* tp_new */
2050 PyObject_GC_Del, /* tp_free */
2051};
2052
2053
2054/* module level code ********************************************************/
2055
2056PyDoc_STRVAR(module_doc,
2057"Functional tools for creating and using iterators.\n\
2058\n\
2059Infinite iterators:\n\
2060count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002061cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002062repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002063\n\
2064Iterators terminating on the shortest input sequence:\n\
2065izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002066ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2067ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002068islice(seq, [start,] stop [, step]) --> elements from\n\
2069 seq[start:stop:step]\n\
2070imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2071starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002072chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002073takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2074dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
2075");
2076
2077
2078PyMODINIT_FUNC
2079inititertools(void)
2080{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002081 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002082 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002083 char *name;
2084 PyTypeObject *typelist[] = {
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00002085 &tee_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002086 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002087 &dropwhile_type,
2088 &takewhile_type,
2089 &islice_type,
2090 &starmap_type,
2091 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002092 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002093 &ifilter_type,
2094 &ifilterfalse_type,
2095 &count_type,
2096 &izip_type,
2097 &repeat_type,
2098 NULL
2099 };
2100
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002101 m = Py_InitModule3("itertools", NULL, module_doc);
2102
Raymond Hettinger60eca932003-02-09 06:40:58 +00002103 for (i=0 ; typelist[i] != NULL ; i++) {
2104 if (PyType_Ready(typelist[i]) < 0)
2105 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002106 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002107 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002108 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002109 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002110 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002111}