blob: 1076ec1f0f6a523391ca9eeec2d4362af7cc88a7 [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
10/* dropwhile object **********************************************************/
11
12typedef struct {
13 PyObject_HEAD
14 PyObject *func;
15 PyObject *it;
16 long start;
17} dropwhileobject;
18
19PyTypeObject dropwhile_type;
20
21static PyObject *
22dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
23{
24 PyObject *func, *seq;
25 PyObject *it;
26 dropwhileobject *lz;
27
28 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
29 return NULL;
30
31 /* Get iterator. */
32 it = PyObject_GetIter(seq);
33 if (it == NULL)
34 return NULL;
35
36 /* create dropwhileobject structure */
37 lz = (dropwhileobject *)type->tp_alloc(type, 0);
38 if (lz == NULL) {
39 Py_DECREF(it);
40 return NULL;
41 }
42 Py_INCREF(func);
43 lz->func = func;
44 lz->it = it;
45 lz->start = 0;
46
47 return (PyObject *)lz;
48}
49
50static void
51dropwhile_dealloc(dropwhileobject *lz)
52{
53 PyObject_GC_UnTrack(lz);
54 Py_XDECREF(lz->func);
55 Py_XDECREF(lz->it);
56 lz->ob_type->tp_free(lz);
57}
58
59static int
60dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
61{
62 int err;
63
64 if (lz->it) {
65 err = visit(lz->it, arg);
66 if (err)
67 return err;
68 }
69 if (lz->func) {
70 err = visit(lz->func, arg);
71 if (err)
72 return err;
73 }
74 return 0;
75}
76
77static PyObject *
78dropwhile_next(dropwhileobject *lz)
79{
80 PyObject *item, *good;
81 long ok;
82
83 for (;;) {
84 item = PyIter_Next(lz->it);
85 if (item == NULL)
86 return NULL;
87 if (lz->start == 1)
88 return item;
89
90 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
91 if (good == NULL) {
92 Py_DECREF(item);
93 return NULL;
94 }
95 ok = PyObject_IsTrue(good);
96 Py_DECREF(good);
97 if (!ok) {
98 lz->start = 1;
99 return item;
100 }
101 Py_DECREF(item);
102 }
103}
104
105static PyObject *
106dropwhile_getiter(PyObject *lz)
107{
108 Py_INCREF(lz);
109 return lz;
110}
111
112PyDoc_STRVAR(dropwhile_doc,
113"dropwhile(predicate, iterable) --> dropwhile object\n\
114\n\
115Drop items from the iterable while predicate(item) is true.\n\
116Afterwards, return every element until the iterable is exhausted.");
117
118PyTypeObject dropwhile_type = {
119 PyObject_HEAD_INIT(NULL)
120 0, /* ob_size */
121 "itertools.dropwhile", /* tp_name */
122 sizeof(dropwhileobject), /* tp_basicsize */
123 0, /* tp_itemsize */
124 /* methods */
125 (destructor)dropwhile_dealloc, /* tp_dealloc */
126 0, /* tp_print */
127 0, /* tp_getattr */
128 0, /* tp_setattr */
129 0, /* tp_compare */
130 0, /* tp_repr */
131 0, /* tp_as_number */
132 0, /* tp_as_sequence */
133 0, /* tp_as_mapping */
134 0, /* tp_hash */
135 0, /* tp_call */
136 0, /* tp_str */
137 PyObject_GenericGetAttr, /* tp_getattro */
138 0, /* tp_setattro */
139 0, /* tp_as_buffer */
140 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
141 Py_TPFLAGS_BASETYPE, /* tp_flags */
142 dropwhile_doc, /* tp_doc */
143 (traverseproc)dropwhile_traverse, /* tp_traverse */
144 0, /* tp_clear */
145 0, /* tp_richcompare */
146 0, /* tp_weaklistoffset */
147 (getiterfunc)dropwhile_getiter, /* tp_iter */
148 (iternextfunc)dropwhile_next, /* tp_iternext */
149 0, /* tp_methods */
150 0, /* tp_members */
151 0, /* tp_getset */
152 0, /* tp_base */
153 0, /* tp_dict */
154 0, /* tp_descr_get */
155 0, /* tp_descr_set */
156 0, /* tp_dictoffset */
157 0, /* tp_init */
158 PyType_GenericAlloc, /* tp_alloc */
159 dropwhile_new, /* tp_new */
160 PyObject_GC_Del, /* tp_free */
161};
162
163
164/* takewhile object **********************************************************/
165
166typedef struct {
167 PyObject_HEAD
168 PyObject *func;
169 PyObject *it;
170 long stop;
171} takewhileobject;
172
173PyTypeObject takewhile_type;
174
175static PyObject *
176takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
177{
178 PyObject *func, *seq;
179 PyObject *it;
180 takewhileobject *lz;
181
182 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
183 return NULL;
184
185 /* Get iterator. */
186 it = PyObject_GetIter(seq);
187 if (it == NULL)
188 return NULL;
189
190 /* create takewhileobject structure */
191 lz = (takewhileobject *)type->tp_alloc(type, 0);
192 if (lz == NULL) {
193 Py_DECREF(it);
194 return NULL;
195 }
196 Py_INCREF(func);
197 lz->func = func;
198 lz->it = it;
199 lz->stop = 0;
200
201 return (PyObject *)lz;
202}
203
204static void
205takewhile_dealloc(takewhileobject *lz)
206{
207 PyObject_GC_UnTrack(lz);
208 Py_XDECREF(lz->func);
209 Py_XDECREF(lz->it);
210 lz->ob_type->tp_free(lz);
211}
212
213static int
214takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
215{
216 int err;
217
218 if (lz->it) {
219 err = visit(lz->it, arg);
220 if (err)
221 return err;
222 }
223 if (lz->func) {
224 err = visit(lz->func, arg);
225 if (err)
226 return err;
227 }
228 return 0;
229}
230
231static PyObject *
232takewhile_next(takewhileobject *lz)
233{
234 PyObject *item, *good;
235 long ok;
236
237 if (lz->stop == 1)
238 return NULL;
239
240 item = PyIter_Next(lz->it);
241 if (item == NULL)
242 return NULL;
243
244 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
245 if (good == NULL) {
246 Py_DECREF(item);
247 return NULL;
248 }
249 ok = PyObject_IsTrue(good);
250 Py_DECREF(good);
251 if (ok)
252 return item;
253 Py_DECREF(item);
254 lz->stop = 1;
255 return NULL;
256}
257
258static PyObject *
259takewhile_getiter(PyObject *lz)
260{
261 Py_INCREF(lz);
262 return lz;
263}
264
265PyDoc_STRVAR(takewhile_doc,
266"takewhile(predicate, iterable) --> takewhile object\n\
267\n\
268Return successive entries from an iterable as long as the \n\
269predicate evaluates to true for each entry.");
270
271PyTypeObject takewhile_type = {
272 PyObject_HEAD_INIT(NULL)
273 0, /* ob_size */
274 "itertools.takewhile", /* tp_name */
275 sizeof(takewhileobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)takewhile_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_compare */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
294 Py_TPFLAGS_BASETYPE, /* tp_flags */
295 takewhile_doc, /* tp_doc */
296 (traverseproc)takewhile_traverse, /* tp_traverse */
297 0, /* tp_clear */
298 0, /* tp_richcompare */
299 0, /* tp_weaklistoffset */
300 (getiterfunc)takewhile_getiter, /* tp_iter */
301 (iternextfunc)takewhile_next, /* tp_iternext */
302 0, /* tp_methods */
303 0, /* tp_members */
304 0, /* tp_getset */
305 0, /* tp_base */
306 0, /* tp_dict */
307 0, /* tp_descr_get */
308 0, /* tp_descr_set */
309 0, /* tp_dictoffset */
310 0, /* tp_init */
311 PyType_GenericAlloc, /* tp_alloc */
312 takewhile_new, /* tp_new */
313 PyObject_GC_Del, /* tp_free */
314};
315
316
317/* islice object ************************************************************/
318
319typedef struct {
320 PyObject_HEAD
321 PyObject *it;
322 long next;
323 long stop;
324 long step;
325 long cnt;
326} isliceobject;
327
328PyTypeObject islice_type;
329
330static PyObject *
331islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
332{
333 PyObject *seq;
334 long a1=0, a2=0, a3=0, start=0, stop=0, step=1;
335 PyObject *it;
336 int numargs;
337 isliceobject *lz;
338
339 numargs = PyTuple_Size(args);
340 if (!PyArg_ParseTuple(args, "Ol|ll:islice", &seq, &a1, &a2, &a3))
341 return NULL;
342
343 if (numargs == 2) {
344 stop = a1;
345 } else if (numargs == 3) {
346 start = a1;
347 stop = a2;
348 } else {
349 start = a1;
350 stop = a2;
351 step = a3;
352 }
353
354 if (start<0 || stop<0) {
355 PyErr_SetString(PyExc_ValueError,
356 "Indices for islice() must be positive.");
357 return NULL;
358 }
359
360 if (step<1) {
361 PyErr_SetString(PyExc_ValueError,
362 "Step must be one or larger for islice().");
363 return NULL;
364 }
365
366 /* Get iterator. */
367 it = PyObject_GetIter(seq);
368 if (it == NULL)
369 return NULL;
370
371 /* create isliceobject structure */
372 lz = (isliceobject *)type->tp_alloc(type, 0);
373 if (lz == NULL) {
374 Py_DECREF(it);
375 return NULL;
376 }
377 lz->it = it;
378 lz->next = start;
379 lz->stop = stop;
380 lz->step = step;
381 lz->cnt = 0L;
382
383 return (PyObject *)lz;
384}
385
386static void
387islice_dealloc(isliceobject *lz)
388{
389 PyObject_GC_UnTrack(lz);
390 Py_XDECREF(lz->it);
391 lz->ob_type->tp_free(lz);
392}
393
394static int
395islice_traverse(isliceobject *lz, visitproc visit, void *arg)
396{
397 if (lz->it)
398 return visit(lz->it, arg);
399 return 0;
400}
401
402static PyObject *
403islice_next(isliceobject *lz)
404{
405 PyObject *item;
406
407 while (lz->cnt < lz->next) {
408 item = PyIter_Next(lz->it);
409 if (item == NULL)
410 return NULL;
411 Py_DECREF(item);
412 lz->cnt++;
413 }
414 if (lz->cnt >= lz->stop)
415 return NULL;
416 item = PyIter_Next(lz->it);
417 if (item == NULL)
418 return NULL;
419 lz->cnt++;
420 lz->next += lz->step;
421 return item;
422}
423
424static PyObject *
425islice_getiter(PyObject *lz)
426{
427 Py_INCREF(lz);
428 return lz;
429}
430
431PyDoc_STRVAR(islice_doc,
432"islice(iterable, [start,] stop [, step]) --> islice object\n\
433\n\
434Return an iterator whose next() method returns selected values from an\n\
435iterable. If start is specified, will skip all preceding elements;\n\
436otherwise, start defaults to zero. Step defaults to one. If\n\
437specified as another value, step determines how many values are \n\
438skipped between successive calls. Works like a slice() on a list\n\
439but returns an iterator.");
440
441PyTypeObject islice_type = {
442 PyObject_HEAD_INIT(NULL)
443 0, /* ob_size */
444 "itertools.islice", /* tp_name */
445 sizeof(isliceobject), /* tp_basicsize */
446 0, /* tp_itemsize */
447 /* methods */
448 (destructor)islice_dealloc, /* tp_dealloc */
449 0, /* tp_print */
450 0, /* tp_getattr */
451 0, /* tp_setattr */
452 0, /* tp_compare */
453 0, /* tp_repr */
454 0, /* tp_as_number */
455 0, /* tp_as_sequence */
456 0, /* tp_as_mapping */
457 0, /* tp_hash */
458 0, /* tp_call */
459 0, /* tp_str */
460 PyObject_GenericGetAttr, /* tp_getattro */
461 0, /* tp_setattro */
462 0, /* tp_as_buffer */
463 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
464 Py_TPFLAGS_BASETYPE, /* tp_flags */
465 islice_doc, /* tp_doc */
466 (traverseproc)islice_traverse, /* tp_traverse */
467 0, /* tp_clear */
468 0, /* tp_richcompare */
469 0, /* tp_weaklistoffset */
470 (getiterfunc)islice_getiter, /* tp_iter */
471 (iternextfunc)islice_next, /* tp_iternext */
472 0, /* tp_methods */
473 0, /* tp_members */
474 0, /* tp_getset */
475 0, /* tp_base */
476 0, /* tp_dict */
477 0, /* tp_descr_get */
478 0, /* tp_descr_set */
479 0, /* tp_dictoffset */
480 0, /* tp_init */
481 PyType_GenericAlloc, /* tp_alloc */
482 islice_new, /* tp_new */
483 PyObject_GC_Del, /* tp_free */
484};
485
486
487/* starmap object ************************************************************/
488
489typedef struct {
490 PyObject_HEAD
491 PyObject *func;
492 PyObject *it;
493} starmapobject;
494
495PyTypeObject starmap_type;
496
497static PyObject *
498starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
499{
500 PyObject *func, *seq;
501 PyObject *it;
502 starmapobject *lz;
503
504 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
505 return NULL;
506
507 /* Get iterator. */
508 it = PyObject_GetIter(seq);
509 if (it == NULL)
510 return NULL;
511
512 /* create starmapobject structure */
513 lz = (starmapobject *)type->tp_alloc(type, 0);
514 if (lz == NULL) {
515 Py_DECREF(it);
516 return NULL;
517 }
518 Py_INCREF(func);
519 lz->func = func;
520 lz->it = it;
521
522 return (PyObject *)lz;
523}
524
525static void
526starmap_dealloc(starmapobject *lz)
527{
528 PyObject_GC_UnTrack(lz);
529 Py_XDECREF(lz->func);
530 Py_XDECREF(lz->it);
531 lz->ob_type->tp_free(lz);
532}
533
534static int
535starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
536{
537 int err;
538
539 if (lz->it) {
540 err = visit(lz->it, arg);
541 if (err)
542 return err;
543 }
544 if (lz->func) {
545 err = visit(lz->func, arg);
546 if (err)
547 return err;
548 }
549 return 0;
550}
551
552static PyObject *
553starmap_next(starmapobject *lz)
554{
555 PyObject *args;
556 PyObject *result;
557
558 args = PyIter_Next(lz->it);
559 if (args == NULL)
560 return NULL;
561 result = PyObject_Call(lz->func, args, NULL);
562 Py_DECREF(args);
563 return result;
564}
565
566static PyObject *
567starmap_getiter(PyObject *lz)
568{
569 Py_INCREF(lz);
570 return lz;
571}
572
573PyDoc_STRVAR(starmap_doc,
574"starmap(function, sequence) --> starmap object\n\
575\n\
576Return an iterator whose values are returned from the function evaluated\n\
577with a argument tuple taken from the given sequence.");
578
579PyTypeObject starmap_type = {
580 PyObject_HEAD_INIT(NULL)
581 0, /* ob_size */
582 "itertools.starmap", /* tp_name */
583 sizeof(starmapobject), /* tp_basicsize */
584 0, /* tp_itemsize */
585 /* methods */
586 (destructor)starmap_dealloc, /* tp_dealloc */
587 0, /* tp_print */
588 0, /* tp_getattr */
589 0, /* tp_setattr */
590 0, /* tp_compare */
591 0, /* tp_repr */
592 0, /* tp_as_number */
593 0, /* tp_as_sequence */
594 0, /* tp_as_mapping */
595 0, /* tp_hash */
596 0, /* tp_call */
597 0, /* tp_str */
598 PyObject_GenericGetAttr, /* tp_getattro */
599 0, /* tp_setattro */
600 0, /* tp_as_buffer */
601 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
602 Py_TPFLAGS_BASETYPE, /* tp_flags */
603 starmap_doc, /* tp_doc */
604 (traverseproc)starmap_traverse, /* tp_traverse */
605 0, /* tp_clear */
606 0, /* tp_richcompare */
607 0, /* tp_weaklistoffset */
608 (getiterfunc)starmap_getiter, /* tp_iter */
609 (iternextfunc)starmap_next, /* tp_iternext */
610 0, /* tp_methods */
611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 PyType_GenericAlloc, /* tp_alloc */
620 starmap_new, /* tp_new */
621 PyObject_GC_Del, /* tp_free */
622};
623
624
625/* imap object ************************************************************/
626
627typedef struct {
628 PyObject_HEAD
629 PyObject *iters;
630 PyObject *argtuple;
631 PyObject *func;
632} imapobject;
633
634PyTypeObject imap_type;
635
636static PyObject *
637imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
638{
639 PyObject *it, *iters, *argtuple, *func;
640 imapobject *lz;
641 int numargs, i;
642
643 numargs = PyTuple_Size(args);
644 if (numargs < 2) {
645 PyErr_SetString(PyExc_TypeError,
646 "imap() must have at least two arguments.");
647 return NULL;
648 }
649
650 iters = PyTuple_New(numargs-1);
651 if (iters == NULL)
652 return NULL;
653
654 argtuple = PyTuple_New(numargs-1);
655 if (argtuple == NULL) {
656 Py_DECREF(iters);
657 return NULL;
658 }
659
660 for (i=1 ; i<numargs ; i++) {
661 /* Get iterator. */
662 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
663 if (it == NULL) {
664 Py_DECREF(argtuple);
665 Py_DECREF(iters);
666 return NULL;
667 }
668 PyTuple_SET_ITEM(iters, i-1, it);
669 Py_INCREF(Py_None);
670 PyTuple_SET_ITEM(argtuple, i-1, Py_None);
671 }
672
673 /* create imapobject structure */
674 lz = (imapobject *)type->tp_alloc(type, 0);
675 if (lz == NULL) {
676 Py_DECREF(argtuple);
677 Py_DECREF(iters);
678 return NULL;
679 }
680 lz->iters = iters;
681 lz->argtuple = argtuple;
682 func = PyTuple_GET_ITEM(args, 0);
683 Py_INCREF(func);
684 lz->func = func;
685
686 return (PyObject *)lz;
687}
688
689static void
690imap_dealloc(imapobject *lz)
691{
692 PyObject_GC_UnTrack(lz);
693 Py_XDECREF(lz->argtuple);
694 Py_XDECREF(lz->iters);
695 Py_XDECREF(lz->func);
696 lz->ob_type->tp_free(lz);
697}
698
699static int
700imap_traverse(imapobject *lz, visitproc visit, void *arg)
701{
702 int err;
703
704 if (lz->iters) {
705 err = visit(lz->iters, arg);
706 if (err)
707 return err;
708 }
709 if (lz->argtuple) {
710 err = visit(lz->argtuple, arg);
711 if (err)
712 return err;
713 }
714 if (lz->func) {
715 err = visit(lz->func, arg);
716 if (err)
717 return err;
718 }
719 return 0;
720}
721
722static PyObject *
723imap_next(imapobject *lz)
724{
725 PyObject *val;
726 PyObject *argtuple=lz->argtuple;
727 int numargs, i;
728
729 numargs = PyTuple_Size(lz->iters);
730 if (lz->func == Py_None) {
731 argtuple = PyTuple_New(numargs);
732 if (argtuple == NULL)
733 return NULL;
734
735 for (i=0 ; i<numargs ; i++) {
736 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
737 if (val == NULL) {
738 Py_DECREF(argtuple);
739 return NULL;
740 }
741 PyTuple_SET_ITEM(argtuple, i, val);
742 }
743 return argtuple;
744 } else {
745 for (i=0 ; i<numargs ; i++) {
746 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
747 if (val == NULL)
748 return NULL;
749 Py_DECREF(PyTuple_GET_ITEM(argtuple, i));
750 PyTuple_SET_ITEM(argtuple, i, val);
751 }
752 return PyObject_Call(lz->func, argtuple, NULL);
753 }
754}
755
756static PyObject *
757imap_getiter(PyObject *lz)
758{
759 Py_INCREF(lz);
760 return lz;
761}
762
763PyDoc_STRVAR(imap_doc,
764"imap(func, *iterables) --> imap object\n\
765\n\
766Make an iterator that computes the function using arguments from\n\
767each of the iterables. Like map() except that it returns\n\
768an iterator instead of a list and that it stops when the shortest\n\
769iterable is exhausted instead of filling in None for shorter\n\
770iterables.");
771
772PyTypeObject imap_type = {
773 PyObject_HEAD_INIT(NULL)
774 0, /* ob_size */
775 "itertools.imap", /* tp_name */
776 sizeof(imapobject), /* tp_basicsize */
777 0, /* tp_itemsize */
778 /* methods */
779 (destructor)imap_dealloc, /* tp_dealloc */
780 0, /* tp_print */
781 0, /* tp_getattr */
782 0, /* tp_setattr */
783 0, /* tp_compare */
784 0, /* tp_repr */
785 0, /* tp_as_number */
786 0, /* tp_as_sequence */
787 0, /* tp_as_mapping */
788 0, /* tp_hash */
789 0, /* tp_call */
790 0, /* tp_str */
791 PyObject_GenericGetAttr, /* tp_getattro */
792 0, /* tp_setattro */
793 0, /* tp_as_buffer */
794 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
795 Py_TPFLAGS_BASETYPE, /* tp_flags */
796 imap_doc, /* tp_doc */
797 (traverseproc)imap_traverse, /* tp_traverse */
798 0, /* tp_clear */
799 0, /* tp_richcompare */
800 0, /* tp_weaklistoffset */
801 (getiterfunc)imap_getiter, /* tp_iter */
802 (iternextfunc)imap_next, /* tp_iternext */
803 0, /* tp_methods */
804 0, /* tp_members */
805 0, /* tp_getset */
806 0, /* tp_base */
807 0, /* tp_dict */
808 0, /* tp_descr_get */
809 0, /* tp_descr_set */
810 0, /* tp_dictoffset */
811 0, /* tp_init */
812 PyType_GenericAlloc, /* tp_alloc */
813 imap_new, /* tp_new */
814 PyObject_GC_Del, /* tp_free */
815};
816
817
818/* times object ************************************************************/
819
820typedef struct {
821 PyObject_HEAD
822 PyObject *obj;
823 long cnt;
824} timesobject;
825
826PyTypeObject times_type;
827
828static PyObject *
829times_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
830{
831 timesobject *lz;
832 PyObject *obj = Py_None;
833 long cnt;
834
835 if (!PyArg_ParseTuple(args, "l|O:times", &cnt, &obj))
836 return NULL;
837
838 if (cnt < 0) {
839 PyErr_SetString(PyExc_ValueError,
840 "count for imap() cannot be negative.");
841 return NULL;
842 }
843
844 /* create timesobject structure */
845 lz = (timesobject *)type->tp_alloc(type, 0);
846 if (lz == NULL)
847 return NULL;
848 lz->cnt = cnt;
849 Py_INCREF(obj);
850 lz->obj = obj;
851
852 return (PyObject *)lz;
853}
854
855static void
856times_dealloc(timesobject *lz)
857{
858 PyObject_GC_UnTrack(lz);
859 Py_XDECREF(lz->obj);
860 lz->ob_type->tp_free(lz);
861}
862
863static PyObject *
864times_next(timesobject *lz)
865{
866 PyObject *obj = lz->obj;
867
868 if (lz->cnt > 0) {
869 lz->cnt--;
870 Py_INCREF(obj);
871 return obj;
872 }
873 return NULL;
874}
875
876static PyObject *
877times_getiter(PyObject *lz)
878{
879 Py_INCREF(lz);
880 return lz;
881}
882
883PyDoc_STRVAR(times_doc,
884"times(n [,obj]) --> times object\n\
885\n\
886Return a times object whose .next() method returns n consecutive\n\
887instances of obj (default is None).");
888
889PyTypeObject times_type = {
890 PyObject_HEAD_INIT(NULL)
891 0, /* ob_size */
892 "itertools.times", /* tp_name */
893 sizeof(timesobject), /* tp_basicsize */
894 0, /* tp_itemsize */
895 /* methods */
896 (destructor)times_dealloc, /* tp_dealloc */
897 0, /* tp_print */
898 0, /* tp_getattr */
899 0, /* tp_setattr */
900 0, /* tp_compare */
901 0, /* tp_repr */
902 0, /* tp_as_number */
903 0, /* tp_as_sequence */
904 0, /* tp_as_mapping */
905 0, /* tp_hash */
906 0, /* tp_call */
907 0, /* tp_str */
908 PyObject_GenericGetAttr, /* tp_getattro */
909 0, /* tp_setattro */
910 0, /* tp_as_buffer */
911 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
912 Py_TPFLAGS_BASETYPE, /* tp_flags */
913 times_doc, /* tp_doc */
914 0, /* tp_traverse */
915 0, /* tp_clear */
916 0, /* tp_richcompare */
917 0, /* tp_weaklistoffset */
918 (getiterfunc)times_getiter, /* tp_iter */
919 (iternextfunc)times_next, /* tp_iternext */
920 0, /* tp_methods */
921 0, /* tp_members */
922 0, /* tp_getset */
923 0, /* tp_base */
924 0, /* tp_dict */
925 0, /* tp_descr_get */
926 0, /* tp_descr_set */
927 0, /* tp_dictoffset */
928 0, /* tp_init */
929 PyType_GenericAlloc, /* tp_alloc */
930 times_new, /* tp_new */
931 PyObject_GC_Del, /* tp_free */
932};
933
934
935/* ifilter object ************************************************************/
936
937typedef struct {
938 PyObject_HEAD
939 PyObject *func;
940 PyObject *it;
941 long invert;
942} ifilterobject;
943
944PyTypeObject ifilter_type;
945
946static PyObject *
947ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
948{
949 PyObject *func, *seq, *invert=NULL;
950 PyObject *it;
951 ifilterobject *lz;
952 long inv=0;
953
954 if (!PyArg_UnpackTuple(args, "ifilter", 2, 3, &func, &seq, &invert))
955 return NULL;
956
957 if (invert != NULL && PyObject_IsTrue(invert))
958 inv = 1;
959
960 /* Get iterator. */
961 it = PyObject_GetIter(seq);
962 if (it == NULL)
963 return NULL;
964
965 /* create ifilterobject structure */
966 lz = (ifilterobject *)type->tp_alloc(type, 0);
967 if (lz == NULL) {
968 Py_DECREF(it);
969 return NULL;
970 }
971 Py_INCREF(func);
972 lz->func = func;
973 lz->it = it;
974 lz->invert = inv;
975
976 return (PyObject *)lz;
977}
978
979static void
980ifilter_dealloc(ifilterobject *lz)
981{
982 PyObject_GC_UnTrack(lz);
983 Py_XDECREF(lz->func);
984 Py_XDECREF(lz->it);
985 lz->ob_type->tp_free(lz);
986}
987
988static int
989ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
990{
991 int err;
992
993 if (lz->it) {
994 err = visit(lz->it, arg);
995 if (err)
996 return err;
997 }
998 if (lz->func) {
999 err = visit(lz->func, arg);
1000 if (err)
1001 return err;
1002 }
1003 return 0;
1004}
1005
1006static PyObject *
1007ifilter_next(ifilterobject *lz)
1008{
1009 PyObject *item;
1010 long ok;
1011
1012 for (;;) {
1013 item = PyIter_Next(lz->it);
1014 if (item == NULL)
1015 return NULL;
1016
1017 if (lz->func == Py_None) {
1018 ok = PyObject_IsTrue(item);
1019 } else {
1020 PyObject *good;
1021 good = PyObject_CallFunctionObjArgs(lz->func,
1022 item, NULL);
1023 if (good == NULL) {
1024 Py_DECREF(item);
1025 return NULL;
1026 }
1027 ok = PyObject_IsTrue(good);
1028 Py_DECREF(good);
1029 }
1030 if (ok ^ lz->invert)
1031 return item;
1032 Py_DECREF(item);
1033 }
1034}
1035
1036static PyObject *
1037ifilter_getiter(PyObject *lz)
1038{
1039 Py_INCREF(lz);
1040 return lz;
1041}
1042
1043PyDoc_STRVAR(ifilter_doc,
1044"ifilter(function or None, sequence [, invert]) --> ifilter object\n\
1045\n\
1046Return those items of sequence for which function(item) is true. If\n\
1047invert is set to True, return items for which function(item) if False.\n\
1048If function is None, return the items that are true (unless invert is set).");
1049
1050PyTypeObject ifilter_type = {
1051 PyObject_HEAD_INIT(NULL)
1052 0, /* ob_size */
1053 "itertools.ifilter", /* tp_name */
1054 sizeof(ifilterobject), /* tp_basicsize */
1055 0, /* tp_itemsize */
1056 /* methods */
1057 (destructor)ifilter_dealloc, /* tp_dealloc */
1058 0, /* tp_print */
1059 0, /* tp_getattr */
1060 0, /* tp_setattr */
1061 0, /* tp_compare */
1062 0, /* tp_repr */
1063 0, /* tp_as_number */
1064 0, /* tp_as_sequence */
1065 0, /* tp_as_mapping */
1066 0, /* tp_hash */
1067 0, /* tp_call */
1068 0, /* tp_str */
1069 PyObject_GenericGetAttr, /* tp_getattro */
1070 0, /* tp_setattro */
1071 0, /* tp_as_buffer */
1072 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1073 Py_TPFLAGS_BASETYPE, /* tp_flags */
1074 ifilter_doc, /* tp_doc */
1075 (traverseproc)ifilter_traverse, /* tp_traverse */
1076 0, /* tp_clear */
1077 0, /* tp_richcompare */
1078 0, /* tp_weaklistoffset */
1079 (getiterfunc)ifilter_getiter, /* tp_iter */
1080 (iternextfunc)ifilter_next, /* tp_iternext */
1081 0, /* tp_methods */
1082 0, /* tp_members */
1083 0, /* tp_getset */
1084 0, /* tp_base */
1085 0, /* tp_dict */
1086 0, /* tp_descr_get */
1087 0, /* tp_descr_set */
1088 0, /* tp_dictoffset */
1089 0, /* tp_init */
1090 PyType_GenericAlloc, /* tp_alloc */
1091 ifilter_new, /* tp_new */
1092 PyObject_GC_Del, /* tp_free */
1093};
1094
1095
1096/* count object ************************************************************/
1097
1098typedef struct {
1099 PyObject_HEAD
1100 long cnt;
1101} countobject;
1102
1103PyTypeObject count_type;
1104
1105static PyObject *
1106count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1107{
1108 countobject *lz;
1109 long cnt = 0;
1110
1111 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1112 return NULL;
1113
1114 /* create countobject structure */
1115 lz = (countobject *)PyObject_New(countobject, &count_type);
1116 if (lz == NULL)
1117 return NULL;
1118 lz->cnt = cnt;
1119
1120 return (PyObject *)lz;
1121}
1122
1123static PyObject *
1124count_next(countobject *lz)
1125{
1126 return PyInt_FromLong(lz->cnt++);
1127}
1128
1129static PyObject *
1130count_getiter(PyObject *lz)
1131{
1132 Py_INCREF(lz);
1133 return lz;
1134}
1135
1136PyDoc_STRVAR(count_doc,
1137"count([firstval]) --> count object\n\
1138\n\
1139Return a count object whose .next() method returns consecutive\n\
1140integers starting from zero or, if specified, from firstval.");
1141
1142PyTypeObject count_type = {
1143 PyObject_HEAD_INIT(NULL)
1144 0, /* ob_size */
1145 "itertools.count", /* tp_name */
1146 sizeof(countobject), /* tp_basicsize */
1147 0, /* tp_itemsize */
1148 /* methods */
1149 (destructor)PyObject_Del, /* tp_dealloc */
1150 0, /* tp_print */
1151 0, /* tp_getattr */
1152 0, /* tp_setattr */
1153 0, /* tp_compare */
1154 0, /* tp_repr */
1155 0, /* tp_as_number */
1156 0, /* tp_as_sequence */
1157 0, /* tp_as_mapping */
1158 0, /* tp_hash */
1159 0, /* tp_call */
1160 0, /* tp_str */
1161 PyObject_GenericGetAttr, /* tp_getattro */
1162 0, /* tp_setattro */
1163 0, /* tp_as_buffer */
1164 Py_TPFLAGS_DEFAULT, /* tp_flags */
1165 count_doc, /* tp_doc */
1166 0, /* tp_traverse */
1167 0, /* tp_clear */
1168 0, /* tp_richcompare */
1169 0, /* tp_weaklistoffset */
1170 (getiterfunc)count_getiter, /* tp_iter */
1171 (iternextfunc)count_next, /* tp_iternext */
1172 0, /* tp_methods */
1173 0, /* tp_members */
1174 0, /* tp_getset */
1175 0, /* tp_base */
1176 0, /* tp_dict */
1177 0, /* tp_descr_get */
1178 0, /* tp_descr_set */
1179 0, /* tp_dictoffset */
1180 0, /* tp_init */
1181 PyType_GenericAlloc, /* tp_alloc */
1182 count_new, /* tp_new */
1183};
1184
1185
1186/* izip object ************************************************************/
1187
1188#include "Python.h"
1189
1190typedef struct {
1191 PyObject_HEAD
1192 long tuplesize;
1193 PyObject *ittuple; /* tuple of iterators */
1194} izipobject;
1195
1196PyTypeObject izip_type;
1197
1198static PyObject *
1199izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1200{
1201 izipobject *lz;
1202 int i;
1203 PyObject *ittuple; /* tuple of iterators */
1204 int tuplesize = PySequence_Length(args);
1205
1206 if (tuplesize < 1) {
1207 PyErr_SetString(PyExc_TypeError,
1208 "izip() requires at least one sequence");
1209 return NULL;
1210 }
1211
1212 /* args must be a tuple */
1213 assert(PyTuple_Check(args));
1214
1215 /* obtain iterators */
1216 ittuple = PyTuple_New(tuplesize);
1217 if(ittuple == NULL)
1218 return NULL;
1219 for (i=0; i < tuplesize; ++i) {
1220 PyObject *item = PyTuple_GET_ITEM(args, i);
1221 PyObject *it = PyObject_GetIter(item);
1222 if (it == NULL) {
1223 if (PyErr_ExceptionMatches(PyExc_TypeError))
1224 PyErr_Format(PyExc_TypeError,
1225 "izip argument #%d must support iteration",
1226 i+1);
1227 Py_DECREF(ittuple);
1228 return NULL;
1229 }
1230 PyTuple_SET_ITEM(ittuple, i, it);
1231 }
1232
1233 /* create izipobject structure */
1234 lz = (izipobject *)type->tp_alloc(type, 0);
1235 if (lz == NULL) {
1236 Py_DECREF(ittuple);
1237 return NULL;
1238 }
1239 lz->ittuple = ittuple;
1240 lz->tuplesize = tuplesize;
1241
1242 return (PyObject *)lz;
1243}
1244
1245static void
1246izip_dealloc(izipobject *lz)
1247{
1248 PyObject_GC_UnTrack(lz);
1249 Py_XDECREF(lz->ittuple);
1250 lz->ob_type->tp_free(lz);
1251}
1252
1253static int
1254izip_traverse(izipobject *lz, visitproc visit, void *arg)
1255{
1256 if (lz->ittuple)
1257 return visit(lz->ittuple, arg);
1258 return 0;
1259}
1260
1261static PyObject *
1262izip_next(izipobject *lz)
1263{
1264 int i;
1265 long tuplesize = lz->tuplesize;
1266 PyObject *result;
1267 PyObject *it;
1268 PyObject *item;
1269
1270 result = PyTuple_New(tuplesize);
1271 if (result == NULL)
1272 return NULL;
1273
1274 for (i=0 ; i < tuplesize ; i++) {
1275 it = PyTuple_GET_ITEM(lz->ittuple, i);
1276 item = PyIter_Next(it);
1277 if (item == NULL) {
1278 Py_DECREF(result);
1279 return NULL;
1280 }
1281 PyTuple_SET_ITEM(result, i, item);
1282 }
1283 return result;
1284}
1285
1286static PyObject *
1287izip_getiter(PyObject *lz)
1288{
1289 Py_INCREF(lz);
1290 return lz;
1291}
1292
1293PyDoc_STRVAR(izip_doc,
1294"izip(iter1 [,iter2 [...]]) --> izip object\n\
1295\n\
1296Return a izip object whose .next() method returns a tuple where\n\
1297the i-th element comes from the i-th iterable argument. The .next()\n\
1298method continues until the shortest iterable in the argument sequence\n\
1299is exhausted and then it raises StopIteration. Works like the zip()\n\
1300function but consumes less memory by returning an iterator instead of\n\
1301a list.");
1302
1303PyTypeObject izip_type = {
1304 PyObject_HEAD_INIT(NULL)
1305 0, /* ob_size */
1306 "itertools.izip", /* tp_name */
1307 sizeof(izipobject), /* tp_basicsize */
1308 0, /* tp_itemsize */
1309 /* methods */
1310 (destructor)izip_dealloc, /* tp_dealloc */
1311 0, /* tp_print */
1312 0, /* tp_getattr */
1313 0, /* tp_setattr */
1314 0, /* tp_compare */
1315 0, /* tp_repr */
1316 0, /* tp_as_number */
1317 0, /* tp_as_sequence */
1318 0, /* tp_as_mapping */
1319 0, /* tp_hash */
1320 0, /* tp_call */
1321 0, /* tp_str */
1322 PyObject_GenericGetAttr, /* tp_getattro */
1323 0, /* tp_setattro */
1324 0, /* tp_as_buffer */
1325 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1326 Py_TPFLAGS_BASETYPE, /* tp_flags */
1327 izip_doc, /* tp_doc */
1328 (traverseproc)izip_traverse, /* tp_traverse */
1329 0, /* tp_clear */
1330 0, /* tp_richcompare */
1331 0, /* tp_weaklistoffset */
1332 (getiterfunc)izip_getiter, /* tp_iter */
1333 (iternextfunc)izip_next, /* tp_iternext */
1334 0, /* tp_methods */
1335 0, /* tp_members */
1336 0, /* tp_getset */
1337 0, /* tp_base */
1338 0, /* tp_dict */
1339 0, /* tp_descr_get */
1340 0, /* tp_descr_set */
1341 0, /* tp_dictoffset */
1342 0, /* tp_init */
1343 PyType_GenericAlloc, /* tp_alloc */
1344 izip_new, /* tp_new */
1345 PyObject_GC_Del, /* tp_free */
1346};
1347
1348
1349/* repeat object ************************************************************/
1350
1351typedef struct {
1352 PyObject_HEAD
1353 PyObject *element;
1354} repeatobject;
1355
1356PyTypeObject repeat_type;
1357
1358static PyObject *
1359repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1360{
1361 repeatobject *ro;
1362 PyObject *element;
1363
1364 if (!PyArg_UnpackTuple(args, "repeat", 1, 1, &element))
1365 return NULL;
1366
1367 ro = (repeatobject *)type->tp_alloc(type, 0);
1368 if (ro == NULL)
1369 return NULL;
1370 Py_INCREF(element);
1371 ro->element = element;
1372 return (PyObject *)ro;
1373}
1374
1375static void
1376repeat_dealloc(repeatobject *ro)
1377{
1378 PyObject_GC_UnTrack(ro);
1379 Py_XDECREF(ro->element);
1380 ro->ob_type->tp_free(ro);
1381}
1382
1383static int
1384repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
1385{
1386 if (ro->element)
1387 return visit(ro->element, arg);
1388 return 0;
1389}
1390
1391static PyObject *
1392repeat_next(repeatobject *ro)
1393{
1394 Py_INCREF(ro->element);
1395 return ro->element;
1396}
1397
1398static PyObject *
1399repeat_getiter(PyObject *ro)
1400{
1401 Py_INCREF(ro);
1402 return ro;
1403}
1404
1405PyDoc_STRVAR(repeat_doc,
1406"repeat(element) -> create an iterator which returns the element endlessly.");
1407
1408PyTypeObject repeat_type = {
1409 PyObject_HEAD_INIT(NULL)
1410 0, /* ob_size */
1411 "itertools.repeat", /* tp_name */
1412 sizeof(repeatobject), /* tp_basicsize */
1413 0, /* tp_itemsize */
1414 /* methods */
1415 (destructor)repeat_dealloc, /* tp_dealloc */
1416 0, /* tp_print */
1417 0, /* tp_getattr */
1418 0, /* tp_setattr */
1419 0, /* tp_compare */
1420 0, /* tp_repr */
1421 0, /* tp_as_number */
1422 0, /* tp_as_sequence */
1423 0, /* tp_as_mapping */
1424 0, /* tp_hash */
1425 0, /* tp_call */
1426 0, /* tp_str */
1427 PyObject_GenericGetAttr, /* tp_getattro */
1428 0, /* tp_setattro */
1429 0, /* tp_as_buffer */
1430 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1431 Py_TPFLAGS_BASETYPE, /* tp_flags */
1432 repeat_doc, /* tp_doc */
1433 (traverseproc)repeat_traverse, /* tp_traverse */
1434 0, /* tp_clear */
1435 0, /* tp_richcompare */
1436 0, /* tp_weaklistoffset */
1437 (getiterfunc)repeat_getiter, /* tp_iter */
1438 (iternextfunc)repeat_next, /* tp_iternext */
1439 0, /* tp_methods */
1440 0, /* tp_members */
1441 0, /* tp_getset */
1442 0, /* tp_base */
1443 0, /* tp_dict */
1444 0, /* tp_descr_get */
1445 0, /* tp_descr_set */
1446 0, /* tp_dictoffset */
1447 0, /* tp_init */
1448 PyType_GenericAlloc, /* tp_alloc */
1449 repeat_new, /* tp_new */
1450 PyObject_GC_Del, /* tp_free */
1451};
1452
1453
1454/* module level code ********************************************************/
1455
1456PyDoc_STRVAR(module_doc,
1457"Functional tools for creating and using iterators.\n\
1458\n\
1459Infinite iterators:\n\
1460count([n]) --> n, n+1, n+2, ...\n\
1461repeat(elem) --> elem, elem, elem, ...\n\
1462\n\
1463Iterators terminating on the shortest input sequence:\n\
1464izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
1465ifilter(pred, seq, invert=False) --> elements of seq where\n\
1466 pred(elem) is True (or False if invert is set)\n\
1467islice(seq, [start,] stop [, step]) --> elements from\n\
1468 seq[start:stop:step]\n\
1469imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
1470starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
1471times(n, [obj]) --> obj, obj, ... for n times. obj defaults to None\n\
1472takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
1473dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
1474");
1475
1476
1477PyMODINIT_FUNC
1478inititertools(void)
1479{
1480 PyObject *m;
1481 m = Py_InitModule3("itertools", NULL, module_doc);
1482
1483 PyModule_AddObject(m, "dropwhile", (PyObject *)&dropwhile_type);
1484 if (PyType_Ready(&dropwhile_type) < 0)
1485 return;
1486 Py_INCREF(&dropwhile_type);
1487
1488 PyModule_AddObject(m, "takewhile", (PyObject *)&takewhile_type);
1489 if (PyType_Ready(&takewhile_type) < 0)
1490 return;
1491 Py_INCREF(&takewhile_type);
1492
1493 PyModule_AddObject(m, "islice", (PyObject *)&islice_type);
1494 if (PyType_Ready(&islice_type) < 0)
1495 return;
1496 Py_INCREF(&islice_type);
1497
1498 PyModule_AddObject(m, "starmap", (PyObject *)&starmap_type);
1499 if (PyType_Ready(&starmap_type) < 0)
1500 return;
1501 Py_INCREF(&starmap_type);
1502
1503 PyModule_AddObject(m, "imap", (PyObject *)&imap_type);
1504 if (PyType_Ready(&imap_type) < 0)
1505 return;
1506 Py_INCREF(&imap_type);
1507
1508 PyModule_AddObject(m, "times", (PyObject *)&times_type);
1509 if (PyType_Ready(&times_type) < 0)
1510 return;
1511 Py_INCREF(&times_type);
1512
1513 if (PyType_Ready(&ifilter_type) < 0)
1514 return;
1515 Py_INCREF(&ifilter_type);
1516 PyModule_AddObject(m, "ifilter", (PyObject *)&ifilter_type);
1517
1518 if (PyType_Ready(&count_type) < 0)
1519 return;
1520 Py_INCREF(&count_type);
1521 PyModule_AddObject(m, "count", (PyObject *)&count_type);
1522
1523 if (PyType_Ready(&izip_type) < 0)
1524 return;
1525 Py_INCREF(&izip_type);
1526 PyModule_AddObject(m, "izip", (PyObject *)&izip_type);
1527
1528 if (PyType_Ready(&repeat_type) < 0)
1529 return;
1530 Py_INCREF(&repeat_type);
1531 PyModule_AddObject(m, "repeat", (PyObject *)&repeat_type);
1532}