blob: fae4511ef7be9d747b5a5f0483d760dce920f7ad [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 Hettinger61fe64d2003-02-23 04:40:07 +000010/* cycle object **********************************************************/
11
12typedef struct {
13 PyObject_HEAD
14 PyObject *it;
15 PyObject *saved;
16 int firstpass;
17} cycleobject;
18
19PyTypeObject cycle_type;
20
21static PyObject *
22cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
23{
24 PyObject *it;
25 PyObject *iterable;
26 PyObject *saved;
27 cycleobject *lz;
28
29 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
30 return NULL;
31
32 /* Get iterator. */
33 it = PyObject_GetIter(iterable);
34 if (it == NULL)
35 return NULL;
36
37 saved = PyList_New(0);
38 if (saved == NULL) {
39 Py_DECREF(it);
40 return NULL;
41 }
42
43 /* create cycleobject structure */
44 lz = (cycleobject *)type->tp_alloc(type, 0);
45 if (lz == NULL) {
46 Py_DECREF(it);
47 Py_DECREF(saved);
48 return NULL;
49 }
50 lz->it = it;
51 lz->saved = saved;
52 lz->firstpass = 0;
53
54 return (PyObject *)lz;
55}
56
57static void
58cycle_dealloc(cycleobject *lz)
59{
60 PyObject_GC_UnTrack(lz);
61 Py_XDECREF(lz->saved);
62 Py_XDECREF(lz->it);
63 lz->ob_type->tp_free(lz);
64}
65
66static int
67cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
68{
69 int err;
70
71 if (lz->it) {
72 err = visit(lz->it, arg);
73 if (err)
74 return err;
75 }
76 if (lz->saved) {
77 err = visit(lz->saved, arg);
78 if (err)
79 return err;
80 }
81 return 0;
82}
83
84static PyObject *
85cycle_next(cycleobject *lz)
86{
87 PyObject *item;
88 PyObject *it;
89
90 while (1) {
91 item = PyIter_Next(lz->it);
92 if (item != NULL) {
93 if (!lz->firstpass)
94 PyList_Append(lz->saved, item);
95 return item;
96 }
97 if (PyList_Size(lz->saved) == 0)
98 return NULL;
99 it = PyObject_GetIter(lz->saved);
100 if (it == NULL)
101 return NULL;
102 Py_DECREF(lz->it);
103 lz->it = it;
104 lz->firstpass = 1;
105 }
106}
107
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000108PyDoc_STRVAR(cycle_doc,
109"cycle(iterable) --> cycle object\n\
110\n\
111Return elements from the iterable until it is exhausted.\n\
112Then repeat the sequence indefinitely.");
113
114PyTypeObject cycle_type = {
115 PyObject_HEAD_INIT(NULL)
116 0, /* ob_size */
117 "itertools.cycle", /* tp_name */
118 sizeof(cycleobject), /* tp_basicsize */
119 0, /* tp_itemsize */
120 /* methods */
121 (destructor)cycle_dealloc, /* tp_dealloc */
122 0, /* tp_print */
123 0, /* tp_getattr */
124 0, /* tp_setattr */
125 0, /* tp_compare */
126 0, /* tp_repr */
127 0, /* tp_as_number */
128 0, /* tp_as_sequence */
129 0, /* tp_as_mapping */
130 0, /* tp_hash */
131 0, /* tp_call */
132 0, /* tp_str */
133 PyObject_GenericGetAttr, /* tp_getattro */
134 0, /* tp_setattro */
135 0, /* tp_as_buffer */
136 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
137 Py_TPFLAGS_BASETYPE, /* tp_flags */
138 cycle_doc, /* tp_doc */
139 (traverseproc)cycle_traverse, /* tp_traverse */
140 0, /* tp_clear */
141 0, /* tp_richcompare */
142 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000143 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000144 (iternextfunc)cycle_next, /* tp_iternext */
145 0, /* tp_methods */
146 0, /* tp_members */
147 0, /* tp_getset */
148 0, /* tp_base */
149 0, /* tp_dict */
150 0, /* tp_descr_get */
151 0, /* tp_descr_set */
152 0, /* tp_dictoffset */
153 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000154 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000155 cycle_new, /* tp_new */
156 PyObject_GC_Del, /* tp_free */
157};
158
159
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000160/* dropwhile object **********************************************************/
161
162typedef struct {
163 PyObject_HEAD
164 PyObject *func;
165 PyObject *it;
166 long start;
167} dropwhileobject;
168
169PyTypeObject dropwhile_type;
170
171static PyObject *
172dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
173{
174 PyObject *func, *seq;
175 PyObject *it;
176 dropwhileobject *lz;
177
178 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
179 return NULL;
180
181 /* Get iterator. */
182 it = PyObject_GetIter(seq);
183 if (it == NULL)
184 return NULL;
185
186 /* create dropwhileobject structure */
187 lz = (dropwhileobject *)type->tp_alloc(type, 0);
188 if (lz == NULL) {
189 Py_DECREF(it);
190 return NULL;
191 }
192 Py_INCREF(func);
193 lz->func = func;
194 lz->it = it;
195 lz->start = 0;
196
197 return (PyObject *)lz;
198}
199
200static void
201dropwhile_dealloc(dropwhileobject *lz)
202{
203 PyObject_GC_UnTrack(lz);
204 Py_XDECREF(lz->func);
205 Py_XDECREF(lz->it);
206 lz->ob_type->tp_free(lz);
207}
208
209static int
210dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
211{
212 int err;
213
214 if (lz->it) {
215 err = visit(lz->it, arg);
216 if (err)
217 return err;
218 }
219 if (lz->func) {
220 err = visit(lz->func, arg);
221 if (err)
222 return err;
223 }
224 return 0;
225}
226
227static PyObject *
228dropwhile_next(dropwhileobject *lz)
229{
230 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000231 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000232 long ok;
233
234 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000235 assert(PyIter_Check(it));
236 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000237 if (item == NULL)
238 return NULL;
239 if (lz->start == 1)
240 return item;
241
242 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
243 if (good == NULL) {
244 Py_DECREF(item);
245 return NULL;
246 }
247 ok = PyObject_IsTrue(good);
248 Py_DECREF(good);
249 if (!ok) {
250 lz->start = 1;
251 return item;
252 }
253 Py_DECREF(item);
254 }
255}
256
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000257PyDoc_STRVAR(dropwhile_doc,
258"dropwhile(predicate, iterable) --> dropwhile object\n\
259\n\
260Drop items from the iterable while predicate(item) is true.\n\
261Afterwards, return every element until the iterable is exhausted.");
262
263PyTypeObject dropwhile_type = {
264 PyObject_HEAD_INIT(NULL)
265 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000266 "itertools.dropwhile", /* tp_name */
267 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000268 0, /* tp_itemsize */
269 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000270 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000271 0, /* tp_print */
272 0, /* tp_getattr */
273 0, /* tp_setattr */
274 0, /* tp_compare */
275 0, /* tp_repr */
276 0, /* tp_as_number */
277 0, /* tp_as_sequence */
278 0, /* tp_as_mapping */
279 0, /* tp_hash */
280 0, /* tp_call */
281 0, /* tp_str */
282 PyObject_GenericGetAttr, /* tp_getattro */
283 0, /* tp_setattro */
284 0, /* tp_as_buffer */
285 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
286 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000287 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000288 (traverseproc)dropwhile_traverse, /* tp_traverse */
289 0, /* tp_clear */
290 0, /* tp_richcompare */
291 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000292 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000293 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000294 0, /* tp_methods */
295 0, /* tp_members */
296 0, /* tp_getset */
297 0, /* tp_base */
298 0, /* tp_dict */
299 0, /* tp_descr_get */
300 0, /* tp_descr_set */
301 0, /* tp_dictoffset */
302 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000303 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000304 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000305 PyObject_GC_Del, /* tp_free */
306};
307
308
309/* takewhile object **********************************************************/
310
311typedef struct {
312 PyObject_HEAD
313 PyObject *func;
314 PyObject *it;
315 long stop;
316} takewhileobject;
317
318PyTypeObject takewhile_type;
319
320static PyObject *
321takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
322{
323 PyObject *func, *seq;
324 PyObject *it;
325 takewhileobject *lz;
326
327 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
328 return NULL;
329
330 /* Get iterator. */
331 it = PyObject_GetIter(seq);
332 if (it == NULL)
333 return NULL;
334
335 /* create takewhileobject structure */
336 lz = (takewhileobject *)type->tp_alloc(type, 0);
337 if (lz == NULL) {
338 Py_DECREF(it);
339 return NULL;
340 }
341 Py_INCREF(func);
342 lz->func = func;
343 lz->it = it;
344 lz->stop = 0;
345
346 return (PyObject *)lz;
347}
348
349static void
350takewhile_dealloc(takewhileobject *lz)
351{
352 PyObject_GC_UnTrack(lz);
353 Py_XDECREF(lz->func);
354 Py_XDECREF(lz->it);
355 lz->ob_type->tp_free(lz);
356}
357
358static int
359takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
360{
361 int err;
362
363 if (lz->it) {
364 err = visit(lz->it, arg);
365 if (err)
366 return err;
367 }
368 if (lz->func) {
369 err = visit(lz->func, arg);
370 if (err)
371 return err;
372 }
373 return 0;
374}
375
376static PyObject *
377takewhile_next(takewhileobject *lz)
378{
379 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000380 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000381 long ok;
382
383 if (lz->stop == 1)
384 return NULL;
385
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000386 assert(PyIter_Check(it));
387 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000388 if (item == NULL)
389 return NULL;
390
391 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
392 if (good == NULL) {
393 Py_DECREF(item);
394 return NULL;
395 }
396 ok = PyObject_IsTrue(good);
397 Py_DECREF(good);
398 if (ok)
399 return item;
400 Py_DECREF(item);
401 lz->stop = 1;
402 return NULL;
403}
404
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000405PyDoc_STRVAR(takewhile_doc,
406"takewhile(predicate, iterable) --> takewhile object\n\
407\n\
408Return successive entries from an iterable as long as the \n\
409predicate evaluates to true for each entry.");
410
411PyTypeObject takewhile_type = {
412 PyObject_HEAD_INIT(NULL)
413 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000414 "itertools.takewhile", /* tp_name */
415 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000416 0, /* tp_itemsize */
417 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000418 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000419 0, /* tp_print */
420 0, /* tp_getattr */
421 0, /* tp_setattr */
422 0, /* tp_compare */
423 0, /* tp_repr */
424 0, /* tp_as_number */
425 0, /* tp_as_sequence */
426 0, /* tp_as_mapping */
427 0, /* tp_hash */
428 0, /* tp_call */
429 0, /* tp_str */
430 PyObject_GenericGetAttr, /* tp_getattro */
431 0, /* tp_setattro */
432 0, /* tp_as_buffer */
433 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
434 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000435 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000436 (traverseproc)takewhile_traverse, /* tp_traverse */
437 0, /* tp_clear */
438 0, /* tp_richcompare */
439 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000440 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000441 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000442 0, /* tp_methods */
443 0, /* tp_members */
444 0, /* tp_getset */
445 0, /* tp_base */
446 0, /* tp_dict */
447 0, /* tp_descr_get */
448 0, /* tp_descr_set */
449 0, /* tp_dictoffset */
450 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000451 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000452 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000453 PyObject_GC_Del, /* tp_free */
454};
455
456
457/* islice object ************************************************************/
458
459typedef struct {
460 PyObject_HEAD
461 PyObject *it;
462 long next;
463 long stop;
464 long step;
465 long cnt;
466} isliceobject;
467
468PyTypeObject islice_type;
469
470static PyObject *
471islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
472{
473 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000474 long start=0, stop=-1, step=1;
475 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000476 int numargs;
477 isliceobject *lz;
478
479 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +0000480 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000481 return NULL;
482
483 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000484 if (a1 != Py_None) {
485 stop = PyInt_AsLong(a1);
486 if (stop == -1) {
487 if (PyErr_Occurred())
488 PyErr_Clear();
489 PyErr_SetString(PyExc_ValueError,
490 "Stop argument must be an integer or None.");
491 return NULL;
492 }
493 }
Raymond Hettinger341deb72003-05-02 19:44:20 +0000494 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000495 start = PyInt_AsLong(a1);
496 if (start == -1 && PyErr_Occurred()) {
497 PyErr_Clear();
498 PyErr_SetString(PyExc_ValueError,
499 "Start argument must be an integer.");
500 return NULL;
501 }
502 if (a2 != Py_None) {
503 stop = PyInt_AsLong(a2);
504 if (stop == -1) {
505 if (PyErr_Occurred())
506 PyErr_Clear();
507 PyErr_SetString(PyExc_ValueError,
508 "Stop argument must be an integer or None.");
509 return NULL;
510 }
511 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000512 }
513
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000514 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000515 PyErr_SetString(PyExc_ValueError,
516 "Indices for islice() must be positive.");
517 return NULL;
518 }
519
520 if (step<1) {
521 PyErr_SetString(PyExc_ValueError,
522 "Step must be one or larger for islice().");
523 return NULL;
524 }
525
526 /* Get iterator. */
527 it = PyObject_GetIter(seq);
528 if (it == NULL)
529 return NULL;
530
531 /* create isliceobject structure */
532 lz = (isliceobject *)type->tp_alloc(type, 0);
533 if (lz == NULL) {
534 Py_DECREF(it);
535 return NULL;
536 }
537 lz->it = it;
538 lz->next = start;
539 lz->stop = stop;
540 lz->step = step;
541 lz->cnt = 0L;
542
543 return (PyObject *)lz;
544}
545
546static void
547islice_dealloc(isliceobject *lz)
548{
549 PyObject_GC_UnTrack(lz);
550 Py_XDECREF(lz->it);
551 lz->ob_type->tp_free(lz);
552}
553
554static int
555islice_traverse(isliceobject *lz, visitproc visit, void *arg)
556{
557 if (lz->it)
558 return visit(lz->it, arg);
559 return 0;
560}
561
562static PyObject *
563islice_next(isliceobject *lz)
564{
565 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000566 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +0000567 long oldnext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000568
569 while (lz->cnt < lz->next) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000570 assert(PyIter_Check(it));
571 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000572 if (item == NULL)
573 return NULL;
574 Py_DECREF(item);
575 lz->cnt++;
576 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000577 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000578 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000579 assert(PyIter_Check(it));
580 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000581 if (item == NULL)
582 return NULL;
583 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +0000584 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000585 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +0000586 if (lz->next < oldnext) /* Check for overflow */
587 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000588 return item;
589}
590
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000591PyDoc_STRVAR(islice_doc,
592"islice(iterable, [start,] stop [, step]) --> islice object\n\
593\n\
594Return an iterator whose next() method returns selected values from an\n\
595iterable. If start is specified, will skip all preceding elements;\n\
596otherwise, start defaults to zero. Step defaults to one. If\n\
597specified as another value, step determines how many values are \n\
598skipped between successive calls. Works like a slice() on a list\n\
599but returns an iterator.");
600
601PyTypeObject islice_type = {
602 PyObject_HEAD_INIT(NULL)
603 0, /* ob_size */
604 "itertools.islice", /* tp_name */
605 sizeof(isliceobject), /* tp_basicsize */
606 0, /* tp_itemsize */
607 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000608 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000609 0, /* tp_print */
610 0, /* tp_getattr */
611 0, /* tp_setattr */
612 0, /* tp_compare */
613 0, /* tp_repr */
614 0, /* tp_as_number */
615 0, /* tp_as_sequence */
616 0, /* tp_as_mapping */
617 0, /* tp_hash */
618 0, /* tp_call */
619 0, /* tp_str */
620 PyObject_GenericGetAttr, /* tp_getattro */
621 0, /* tp_setattro */
622 0, /* tp_as_buffer */
623 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
624 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000625 islice_doc, /* tp_doc */
626 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000627 0, /* tp_clear */
628 0, /* tp_richcompare */
629 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000630 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000631 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000632 0, /* tp_methods */
633 0, /* tp_members */
634 0, /* tp_getset */
635 0, /* tp_base */
636 0, /* tp_dict */
637 0, /* tp_descr_get */
638 0, /* tp_descr_set */
639 0, /* tp_dictoffset */
640 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000641 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000642 islice_new, /* tp_new */
643 PyObject_GC_Del, /* tp_free */
644};
645
646
647/* starmap object ************************************************************/
648
649typedef struct {
650 PyObject_HEAD
651 PyObject *func;
652 PyObject *it;
653} starmapobject;
654
655PyTypeObject starmap_type;
656
657static PyObject *
658starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
659{
660 PyObject *func, *seq;
661 PyObject *it;
662 starmapobject *lz;
663
664 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
665 return NULL;
666
667 /* Get iterator. */
668 it = PyObject_GetIter(seq);
669 if (it == NULL)
670 return NULL;
671
672 /* create starmapobject structure */
673 lz = (starmapobject *)type->tp_alloc(type, 0);
674 if (lz == NULL) {
675 Py_DECREF(it);
676 return NULL;
677 }
678 Py_INCREF(func);
679 lz->func = func;
680 lz->it = it;
681
682 return (PyObject *)lz;
683}
684
685static void
686starmap_dealloc(starmapobject *lz)
687{
688 PyObject_GC_UnTrack(lz);
689 Py_XDECREF(lz->func);
690 Py_XDECREF(lz->it);
691 lz->ob_type->tp_free(lz);
692}
693
694static int
695starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
696{
697 int err;
698
699 if (lz->it) {
700 err = visit(lz->it, arg);
701 if (err)
702 return err;
703 }
704 if (lz->func) {
705 err = visit(lz->func, arg);
706 if (err)
707 return err;
708 }
709 return 0;
710}
711
712static PyObject *
713starmap_next(starmapobject *lz)
714{
715 PyObject *args;
716 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000717 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000718
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000719 assert(PyIter_Check(it));
720 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000721 if (args == NULL)
722 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +0000723 if (!PyTuple_CheckExact(args)) {
724 Py_DECREF(args);
725 PyErr_SetString(PyExc_TypeError,
726 "iterator must return a tuple");
727 return NULL;
728 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000729 result = PyObject_Call(lz->func, args, NULL);
730 Py_DECREF(args);
731 return result;
732}
733
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000734PyDoc_STRVAR(starmap_doc,
735"starmap(function, sequence) --> starmap object\n\
736\n\
737Return an iterator whose values are returned from the function evaluated\n\
738with a argument tuple taken from the given sequence.");
739
740PyTypeObject starmap_type = {
741 PyObject_HEAD_INIT(NULL)
742 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000743 "itertools.starmap", /* tp_name */
744 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000745 0, /* tp_itemsize */
746 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000747 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000748 0, /* tp_print */
749 0, /* tp_getattr */
750 0, /* tp_setattr */
751 0, /* tp_compare */
752 0, /* tp_repr */
753 0, /* tp_as_number */
754 0, /* tp_as_sequence */
755 0, /* tp_as_mapping */
756 0, /* tp_hash */
757 0, /* tp_call */
758 0, /* tp_str */
759 PyObject_GenericGetAttr, /* tp_getattro */
760 0, /* tp_setattro */
761 0, /* tp_as_buffer */
762 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
763 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000764 starmap_doc, /* tp_doc */
765 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000766 0, /* tp_clear */
767 0, /* tp_richcompare */
768 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000769 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000770 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000771 0, /* tp_methods */
772 0, /* tp_members */
773 0, /* tp_getset */
774 0, /* tp_base */
775 0, /* tp_dict */
776 0, /* tp_descr_get */
777 0, /* tp_descr_set */
778 0, /* tp_dictoffset */
779 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000780 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000781 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000782 PyObject_GC_Del, /* tp_free */
783};
784
785
786/* imap object ************************************************************/
787
788typedef struct {
789 PyObject_HEAD
790 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000791 PyObject *func;
792} imapobject;
793
794PyTypeObject imap_type;
795
796static PyObject *
797imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
798{
Raymond Hettingerf0c00242003-02-07 07:26:25 +0000799 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000800 imapobject *lz;
801 int numargs, i;
802
803 numargs = PyTuple_Size(args);
804 if (numargs < 2) {
805 PyErr_SetString(PyExc_TypeError,
806 "imap() must have at least two arguments.");
807 return NULL;
808 }
809
810 iters = PyTuple_New(numargs-1);
811 if (iters == NULL)
812 return NULL;
813
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000814 for (i=1 ; i<numargs ; i++) {
815 /* Get iterator. */
816 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
817 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000818 Py_DECREF(iters);
819 return NULL;
820 }
821 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000822 }
823
824 /* create imapobject structure */
825 lz = (imapobject *)type->tp_alloc(type, 0);
826 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000827 Py_DECREF(iters);
828 return NULL;
829 }
830 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831 func = PyTuple_GET_ITEM(args, 0);
832 Py_INCREF(func);
833 lz->func = func;
834
835 return (PyObject *)lz;
836}
837
838static void
839imap_dealloc(imapobject *lz)
840{
841 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 Py_XDECREF(lz->iters);
843 Py_XDECREF(lz->func);
844 lz->ob_type->tp_free(lz);
845}
846
847static int
848imap_traverse(imapobject *lz, visitproc visit, void *arg)
849{
850 int err;
851
852 if (lz->iters) {
853 err = visit(lz->iters, arg);
854 if (err)
855 return err;
856 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000857 if (lz->func) {
858 err = visit(lz->func, arg);
859 if (err)
860 return err;
861 }
862 return 0;
863}
864
Raymond Hettinger2012f172003-02-07 05:32:58 +0000865/*
866imap() is an iterator version of __builtins__.map() except that it does
867not have the None fill-in feature. That was intentionally left out for
868the following reasons:
869
870 1) Itertools are designed to be easily combined and chained together.
871 Having all tools stop with the shortest input is a unifying principle
872 that makes it easier to combine finite iterators (supplying data) with
873 infinite iterators like count() and repeat() (for supplying sequential
874 or constant arguments to a function).
875
876 2) In typical use cases for combining itertools, having one finite data
877 supplier run out before another is likely to be an error condition which
878 should not pass silently by automatically supplying None.
879
880 3) The use cases for automatic None fill-in are rare -- not many functions
881 do something useful when a parameter suddenly switches type and becomes
882 None.
883
884 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000885 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +0000886
887 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
888*/
889
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890static PyObject *
891imap_next(imapobject *lz)
892{
893 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +0000894 PyObject *argtuple;
895 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000896 int numargs, i;
897
898 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +0000899 argtuple = PyTuple_New(numargs);
900 if (argtuple == NULL)
901 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000902
Raymond Hettingerf0c00242003-02-07 07:26:25 +0000903 for (i=0 ; i<numargs ; i++) {
904 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
905 if (val == NULL) {
906 Py_DECREF(argtuple);
907 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000908 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +0000909 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000910 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +0000911 if (lz->func == Py_None)
912 return argtuple;
913 result = PyObject_Call(lz->func, argtuple, NULL);
914 Py_DECREF(argtuple);
915 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000916}
917
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000918PyDoc_STRVAR(imap_doc,
919"imap(func, *iterables) --> imap object\n\
920\n\
921Make an iterator that computes the function using arguments from\n\
922each of the iterables. Like map() except that it returns\n\
923an iterator instead of a list and that it stops when the shortest\n\
924iterable is exhausted instead of filling in None for shorter\n\
925iterables.");
926
927PyTypeObject imap_type = {
928 PyObject_HEAD_INIT(NULL)
929 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000930 "itertools.imap", /* tp_name */
931 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000932 0, /* tp_itemsize */
933 /* methods */
934 (destructor)imap_dealloc, /* tp_dealloc */
935 0, /* tp_print */
936 0, /* tp_getattr */
937 0, /* tp_setattr */
938 0, /* tp_compare */
939 0, /* tp_repr */
940 0, /* tp_as_number */
941 0, /* tp_as_sequence */
942 0, /* tp_as_mapping */
943 0, /* tp_hash */
944 0, /* tp_call */
945 0, /* tp_str */
946 PyObject_GenericGetAttr, /* tp_getattro */
947 0, /* tp_setattro */
948 0, /* tp_as_buffer */
949 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
950 Py_TPFLAGS_BASETYPE, /* tp_flags */
951 imap_doc, /* tp_doc */
952 (traverseproc)imap_traverse, /* tp_traverse */
953 0, /* tp_clear */
954 0, /* tp_richcompare */
955 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000956 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000957 (iternextfunc)imap_next, /* tp_iternext */
958 0, /* tp_methods */
959 0, /* tp_members */
960 0, /* tp_getset */
961 0, /* tp_base */
962 0, /* tp_dict */
963 0, /* tp_descr_get */
964 0, /* tp_descr_set */
965 0, /* tp_dictoffset */
966 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000967 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000968 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000969 PyObject_GC_Del, /* tp_free */
970};
971
972
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000973/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000974
975typedef struct {
976 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000977 long tuplesize;
978 long iternum; /* which iterator is active */
979 PyObject *ittuple; /* tuple of iterators */
980} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000981
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000982PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000983
984static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000985chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000986{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000987 chainobject *lz;
988 int tuplesize = PySequence_Length(args);
989 int i;
990 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000991
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000992 /* obtain iterators */
993 assert(PyTuple_Check(args));
994 ittuple = PyTuple_New(tuplesize);
995 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000996 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000997 for (i=0; i < tuplesize; ++i) {
998 PyObject *item = PyTuple_GET_ITEM(args, i);
999 PyObject *it = PyObject_GetIter(item);
1000 if (it == NULL) {
1001 if (PyErr_ExceptionMatches(PyExc_TypeError))
1002 PyErr_Format(PyExc_TypeError,
1003 "chain argument #%d must support iteration",
1004 i+1);
1005 Py_DECREF(ittuple);
1006 return NULL;
1007 }
1008 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001009 }
1010
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001011 /* create chainobject structure */
1012 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001013 if (lz == NULL) {
1014 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001016 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001017
1018 lz->ittuple = ittuple;
1019 lz->iternum = 0;
1020 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001021
1022 return (PyObject *)lz;
1023}
1024
1025static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001026chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001027{
1028 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001029 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001030 lz->ob_type->tp_free(lz);
1031}
1032
Raymond Hettinger2012f172003-02-07 05:32:58 +00001033static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001034chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001035{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001036 if (lz->ittuple)
1037 return visit(lz->ittuple, arg);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001038 return 0;
1039}
1040
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001041static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001042chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001043{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001044 PyObject *it;
1045 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001046
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001047 while (lz->iternum < lz->tuplesize) {
1048 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1049 item = PyIter_Next(it);
1050 if (item != NULL)
1051 return item;
1052 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053 }
1054 return NULL;
1055}
1056
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001057PyDoc_STRVAR(chain_doc,
1058"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001059\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001060Return a chain object whose .next() method returns elements from the\n\
1061first iterable until it is exhausted, then elements from the next\n\
1062iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001063
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001064PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001065 PyObject_HEAD_INIT(NULL)
1066 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001067 "itertools.chain", /* tp_name */
1068 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001069 0, /* tp_itemsize */
1070 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001071 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072 0, /* tp_print */
1073 0, /* tp_getattr */
1074 0, /* tp_setattr */
1075 0, /* tp_compare */
1076 0, /* tp_repr */
1077 0, /* tp_as_number */
1078 0, /* tp_as_sequence */
1079 0, /* tp_as_mapping */
1080 0, /* tp_hash */
1081 0, /* tp_call */
1082 0, /* tp_str */
1083 PyObject_GenericGetAttr, /* tp_getattro */
1084 0, /* tp_setattro */
1085 0, /* tp_as_buffer */
1086 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1087 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001088 chain_doc, /* tp_doc */
1089 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001090 0, /* tp_clear */
1091 0, /* tp_richcompare */
1092 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001093 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001094 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001095 0, /* tp_methods */
1096 0, /* tp_members */
1097 0, /* tp_getset */
1098 0, /* tp_base */
1099 0, /* tp_dict */
1100 0, /* tp_descr_get */
1101 0, /* tp_descr_set */
1102 0, /* tp_dictoffset */
1103 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001104 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001105 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001106 PyObject_GC_Del, /* tp_free */
1107};
1108
1109
1110/* ifilter object ************************************************************/
1111
1112typedef struct {
1113 PyObject_HEAD
1114 PyObject *func;
1115 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116} ifilterobject;
1117
1118PyTypeObject ifilter_type;
1119
1120static PyObject *
1121ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1122{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001123 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 PyObject *it;
1125 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001126
Raymond Hettinger60eca932003-02-09 06:40:58 +00001127 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128 return NULL;
1129
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 /* Get iterator. */
1131 it = PyObject_GetIter(seq);
1132 if (it == NULL)
1133 return NULL;
1134
1135 /* create ifilterobject structure */
1136 lz = (ifilterobject *)type->tp_alloc(type, 0);
1137 if (lz == NULL) {
1138 Py_DECREF(it);
1139 return NULL;
1140 }
1141 Py_INCREF(func);
1142 lz->func = func;
1143 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001144
1145 return (PyObject *)lz;
1146}
1147
1148static void
1149ifilter_dealloc(ifilterobject *lz)
1150{
1151 PyObject_GC_UnTrack(lz);
1152 Py_XDECREF(lz->func);
1153 Py_XDECREF(lz->it);
1154 lz->ob_type->tp_free(lz);
1155}
1156
1157static int
1158ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1159{
1160 int err;
1161
1162 if (lz->it) {
1163 err = visit(lz->it, arg);
1164 if (err)
1165 return err;
1166 }
1167 if (lz->func) {
1168 err = visit(lz->func, arg);
1169 if (err)
1170 return err;
1171 }
1172 return 0;
1173}
1174
1175static PyObject *
1176ifilter_next(ifilterobject *lz)
1177{
1178 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001179 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001180 long ok;
1181
1182 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001183 assert(PyIter_Check(it));
1184 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001185 if (item == NULL)
1186 return NULL;
1187
1188 if (lz->func == Py_None) {
1189 ok = PyObject_IsTrue(item);
1190 } else {
1191 PyObject *good;
1192 good = PyObject_CallFunctionObjArgs(lz->func,
1193 item, NULL);
1194 if (good == NULL) {
1195 Py_DECREF(item);
1196 return NULL;
1197 }
1198 ok = PyObject_IsTrue(good);
1199 Py_DECREF(good);
1200 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001201 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202 return item;
1203 Py_DECREF(item);
1204 }
1205}
1206
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001208"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001209\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001210Return those items of sequence for which function(item) is true.\n\
1211If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001212
1213PyTypeObject ifilter_type = {
1214 PyObject_HEAD_INIT(NULL)
1215 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001216 "itertools.ifilter", /* tp_name */
1217 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218 0, /* tp_itemsize */
1219 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001220 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221 0, /* tp_print */
1222 0, /* tp_getattr */
1223 0, /* tp_setattr */
1224 0, /* tp_compare */
1225 0, /* tp_repr */
1226 0, /* tp_as_number */
1227 0, /* tp_as_sequence */
1228 0, /* tp_as_mapping */
1229 0, /* tp_hash */
1230 0, /* tp_call */
1231 0, /* tp_str */
1232 PyObject_GenericGetAttr, /* tp_getattro */
1233 0, /* tp_setattro */
1234 0, /* tp_as_buffer */
1235 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1236 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001237 ifilter_doc, /* tp_doc */
1238 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001239 0, /* tp_clear */
1240 0, /* tp_richcompare */
1241 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001242 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001243 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001244 0, /* tp_methods */
1245 0, /* tp_members */
1246 0, /* tp_getset */
1247 0, /* tp_base */
1248 0, /* tp_dict */
1249 0, /* tp_descr_get */
1250 0, /* tp_descr_set */
1251 0, /* tp_dictoffset */
1252 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001253 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001254 ifilter_new, /* tp_new */
1255 PyObject_GC_Del, /* tp_free */
1256};
1257
1258
1259/* ifilterfalse object ************************************************************/
1260
1261typedef struct {
1262 PyObject_HEAD
1263 PyObject *func;
1264 PyObject *it;
1265} ifilterfalseobject;
1266
1267PyTypeObject ifilterfalse_type;
1268
1269static PyObject *
1270ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1271{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001272 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001273 PyObject *it;
1274 ifilterfalseobject *lz;
1275
1276 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1277 return NULL;
1278
1279 /* Get iterator. */
1280 it = PyObject_GetIter(seq);
1281 if (it == NULL)
1282 return NULL;
1283
1284 /* create ifilterfalseobject structure */
1285 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1286 if (lz == NULL) {
1287 Py_DECREF(it);
1288 return NULL;
1289 }
1290 Py_INCREF(func);
1291 lz->func = func;
1292 lz->it = it;
1293
1294 return (PyObject *)lz;
1295}
1296
1297static void
1298ifilterfalse_dealloc(ifilterfalseobject *lz)
1299{
1300 PyObject_GC_UnTrack(lz);
1301 Py_XDECREF(lz->func);
1302 Py_XDECREF(lz->it);
1303 lz->ob_type->tp_free(lz);
1304}
1305
1306static int
1307ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1308{
1309 int err;
1310
1311 if (lz->it) {
1312 err = visit(lz->it, arg);
1313 if (err)
1314 return err;
1315 }
1316 if (lz->func) {
1317 err = visit(lz->func, arg);
1318 if (err)
1319 return err;
1320 }
1321 return 0;
1322}
1323
1324static PyObject *
1325ifilterfalse_next(ifilterfalseobject *lz)
1326{
1327 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001328 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001329 long ok;
1330
1331 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001332 assert(PyIter_Check(it));
1333 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001334 if (item == NULL)
1335 return NULL;
1336
1337 if (lz->func == Py_None) {
1338 ok = PyObject_IsTrue(item);
1339 } else {
1340 PyObject *good;
1341 good = PyObject_CallFunctionObjArgs(lz->func,
1342 item, NULL);
1343 if (good == NULL) {
1344 Py_DECREF(item);
1345 return NULL;
1346 }
1347 ok = PyObject_IsTrue(good);
1348 Py_DECREF(good);
1349 }
1350 if (!ok)
1351 return item;
1352 Py_DECREF(item);
1353 }
1354}
1355
Raymond Hettinger60eca932003-02-09 06:40:58 +00001356PyDoc_STRVAR(ifilterfalse_doc,
1357"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1358\n\
1359Return those items of sequence for which function(item) is false.\n\
1360If function is None, return the items that are false.");
1361
1362PyTypeObject ifilterfalse_type = {
1363 PyObject_HEAD_INIT(NULL)
1364 0, /* ob_size */
1365 "itertools.ifilterfalse", /* tp_name */
1366 sizeof(ifilterfalseobject), /* tp_basicsize */
1367 0, /* tp_itemsize */
1368 /* methods */
1369 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1370 0, /* tp_print */
1371 0, /* tp_getattr */
1372 0, /* tp_setattr */
1373 0, /* tp_compare */
1374 0, /* tp_repr */
1375 0, /* tp_as_number */
1376 0, /* tp_as_sequence */
1377 0, /* tp_as_mapping */
1378 0, /* tp_hash */
1379 0, /* tp_call */
1380 0, /* tp_str */
1381 PyObject_GenericGetAttr, /* tp_getattro */
1382 0, /* tp_setattro */
1383 0, /* tp_as_buffer */
1384 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1385 Py_TPFLAGS_BASETYPE, /* tp_flags */
1386 ifilterfalse_doc, /* tp_doc */
1387 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1388 0, /* tp_clear */
1389 0, /* tp_richcompare */
1390 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001391 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001392 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1393 0, /* tp_methods */
1394 0, /* tp_members */
1395 0, /* tp_getset */
1396 0, /* tp_base */
1397 0, /* tp_dict */
1398 0, /* tp_descr_get */
1399 0, /* tp_descr_set */
1400 0, /* tp_dictoffset */
1401 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001402 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001403 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404 PyObject_GC_Del, /* tp_free */
1405};
1406
1407
1408/* count object ************************************************************/
1409
1410typedef struct {
1411 PyObject_HEAD
1412 long cnt;
1413} countobject;
1414
1415PyTypeObject count_type;
1416
1417static PyObject *
1418count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1419{
1420 countobject *lz;
1421 long cnt = 0;
1422
1423 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1424 return NULL;
1425
1426 /* create countobject structure */
1427 lz = (countobject *)PyObject_New(countobject, &count_type);
1428 if (lz == NULL)
1429 return NULL;
1430 lz->cnt = cnt;
1431
1432 return (PyObject *)lz;
1433}
1434
1435static PyObject *
1436count_next(countobject *lz)
1437{
1438 return PyInt_FromLong(lz->cnt++);
1439}
1440
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001441PyDoc_STRVAR(count_doc,
1442"count([firstval]) --> count object\n\
1443\n\
1444Return a count object whose .next() method returns consecutive\n\
1445integers starting from zero or, if specified, from firstval.");
1446
1447PyTypeObject count_type = {
1448 PyObject_HEAD_INIT(NULL)
1449 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001450 "itertools.count", /* tp_name */
1451 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452 0, /* tp_itemsize */
1453 /* methods */
1454 (destructor)PyObject_Del, /* tp_dealloc */
1455 0, /* tp_print */
1456 0, /* tp_getattr */
1457 0, /* tp_setattr */
1458 0, /* tp_compare */
1459 0, /* tp_repr */
1460 0, /* tp_as_number */
1461 0, /* tp_as_sequence */
1462 0, /* tp_as_mapping */
1463 0, /* tp_hash */
1464 0, /* tp_call */
1465 0, /* tp_str */
1466 PyObject_GenericGetAttr, /* tp_getattro */
1467 0, /* tp_setattro */
1468 0, /* tp_as_buffer */
1469 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001470 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471 0, /* tp_traverse */
1472 0, /* tp_clear */
1473 0, /* tp_richcompare */
1474 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001475 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001476 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477 0, /* tp_methods */
1478 0, /* tp_members */
1479 0, /* tp_getset */
1480 0, /* tp_base */
1481 0, /* tp_dict */
1482 0, /* tp_descr_get */
1483 0, /* tp_descr_set */
1484 0, /* tp_dictoffset */
1485 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001486 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487 count_new, /* tp_new */
1488};
1489
1490
1491/* izip object ************************************************************/
1492
1493#include "Python.h"
1494
1495typedef struct {
1496 PyObject_HEAD
1497 long tuplesize;
1498 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00001499 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500} izipobject;
1501
1502PyTypeObject izip_type;
1503
1504static PyObject *
1505izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1506{
1507 izipobject *lz;
1508 int i;
1509 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00001510 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001511 int tuplesize = PySequence_Length(args);
1512
1513 if (tuplesize < 1) {
1514 PyErr_SetString(PyExc_TypeError,
1515 "izip() requires at least one sequence");
1516 return NULL;
1517 }
1518
1519 /* args must be a tuple */
1520 assert(PyTuple_Check(args));
1521
1522 /* obtain iterators */
1523 ittuple = PyTuple_New(tuplesize);
1524 if(ittuple == NULL)
1525 return NULL;
1526 for (i=0; i < tuplesize; ++i) {
1527 PyObject *item = PyTuple_GET_ITEM(args, i);
1528 PyObject *it = PyObject_GetIter(item);
1529 if (it == NULL) {
1530 if (PyErr_ExceptionMatches(PyExc_TypeError))
1531 PyErr_Format(PyExc_TypeError,
1532 "izip argument #%d must support iteration",
1533 i+1);
1534 Py_DECREF(ittuple);
1535 return NULL;
1536 }
1537 PyTuple_SET_ITEM(ittuple, i, it);
1538 }
1539
Raymond Hettinger2012f172003-02-07 05:32:58 +00001540 /* create a result holder */
1541 result = PyTuple_New(tuplesize);
1542 if (result == NULL) {
1543 Py_DECREF(ittuple);
1544 return NULL;
1545 }
1546 for (i=0 ; i < tuplesize ; i++) {
1547 Py_INCREF(Py_None);
1548 PyTuple_SET_ITEM(result, i, Py_None);
1549 }
1550
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551 /* create izipobject structure */
1552 lz = (izipobject *)type->tp_alloc(type, 0);
1553 if (lz == NULL) {
1554 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001555 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556 return NULL;
1557 }
1558 lz->ittuple = ittuple;
1559 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001560 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001561
1562 return (PyObject *)lz;
1563}
1564
1565static void
1566izip_dealloc(izipobject *lz)
1567{
1568 PyObject_GC_UnTrack(lz);
1569 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001570 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001571 lz->ob_type->tp_free(lz);
1572}
1573
1574static int
1575izip_traverse(izipobject *lz, visitproc visit, void *arg)
1576{
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001577 int err;
1578
1579 if (lz->ittuple) {
1580 err = visit(lz->ittuple, arg);
1581 if (err)
1582 return err;
1583 }
1584 if (lz->result) {
1585 err = visit(lz->result, arg);
1586 if (err)
1587 return err;
1588 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589 return 0;
1590}
1591
1592static PyObject *
1593izip_next(izipobject *lz)
1594{
1595 int i;
1596 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001597 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001598 PyObject *it;
1599 PyObject *item;
1600
Raymond Hettinger2012f172003-02-07 05:32:58 +00001601 if (result->ob_refcnt == 1) {
1602 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001603 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001604 assert(PyIter_Check(it));
1605 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001606 if (item == NULL)
1607 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001608 Py_DECREF(PyTuple_GET_ITEM(result, i));
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001609 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001610 }
1611 Py_INCREF(result);
1612 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00001613 result = PyTuple_New(tuplesize);
1614 if (result == NULL)
1615 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001616 for (i=0 ; i < tuplesize ; i++) {
1617 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001618 assert(PyIter_Check(it));
1619 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001620 if (item == NULL) {
1621 Py_DECREF(result);
1622 return NULL;
1623 }
1624 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001625 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001626 }
1627 return result;
1628}
1629
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001630PyDoc_STRVAR(izip_doc,
1631"izip(iter1 [,iter2 [...]]) --> izip object\n\
1632\n\
1633Return a izip object whose .next() method returns a tuple where\n\
1634the i-th element comes from the i-th iterable argument. The .next()\n\
1635method continues until the shortest iterable in the argument sequence\n\
1636is exhausted and then it raises StopIteration. Works like the zip()\n\
1637function but consumes less memory by returning an iterator instead of\n\
1638a list.");
1639
1640PyTypeObject izip_type = {
1641 PyObject_HEAD_INIT(NULL)
1642 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001643 "itertools.izip", /* tp_name */
1644 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645 0, /* tp_itemsize */
1646 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001647 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648 0, /* tp_print */
1649 0, /* tp_getattr */
1650 0, /* tp_setattr */
1651 0, /* tp_compare */
1652 0, /* tp_repr */
1653 0, /* tp_as_number */
1654 0, /* tp_as_sequence */
1655 0, /* tp_as_mapping */
1656 0, /* tp_hash */
1657 0, /* tp_call */
1658 0, /* tp_str */
1659 PyObject_GenericGetAttr, /* tp_getattro */
1660 0, /* tp_setattro */
1661 0, /* tp_as_buffer */
1662 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1663 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001664 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665 (traverseproc)izip_traverse, /* tp_traverse */
1666 0, /* tp_clear */
1667 0, /* tp_richcompare */
1668 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001669 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001670 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671 0, /* tp_methods */
1672 0, /* tp_members */
1673 0, /* tp_getset */
1674 0, /* tp_base */
1675 0, /* tp_dict */
1676 0, /* tp_descr_get */
1677 0, /* tp_descr_set */
1678 0, /* tp_dictoffset */
1679 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001680 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001681 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682 PyObject_GC_Del, /* tp_free */
1683};
1684
1685
1686/* repeat object ************************************************************/
1687
1688typedef struct {
1689 PyObject_HEAD
1690 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001691 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001692} repeatobject;
1693
1694PyTypeObject repeat_type;
1695
1696static PyObject *
1697repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1698{
1699 repeatobject *ro;
1700 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001701 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001702
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001703 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001704 return NULL;
1705
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00001706 if (PyTuple_Size(args) == 2 && cnt < 0)
1707 cnt = 0;
1708
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709 ro = (repeatobject *)type->tp_alloc(type, 0);
1710 if (ro == NULL)
1711 return NULL;
1712 Py_INCREF(element);
1713 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001714 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001715 return (PyObject *)ro;
1716}
1717
1718static void
1719repeat_dealloc(repeatobject *ro)
1720{
1721 PyObject_GC_UnTrack(ro);
1722 Py_XDECREF(ro->element);
1723 ro->ob_type->tp_free(ro);
1724}
1725
1726static int
1727repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
1728{
1729 if (ro->element)
1730 return visit(ro->element, arg);
1731 return 0;
1732}
1733
1734static PyObject *
1735repeat_next(repeatobject *ro)
1736{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001737 if (ro->cnt == 0)
1738 return NULL;
1739 if (ro->cnt > 0)
1740 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001741 Py_INCREF(ro->element);
1742 return ro->element;
1743}
1744
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001746"repeat(element [,times]) -> create an iterator which returns the element\n\
1747for the specified number of times. If not specified, returns the element\n\
1748endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749
1750PyTypeObject repeat_type = {
1751 PyObject_HEAD_INIT(NULL)
1752 0, /* ob_size */
1753 "itertools.repeat", /* tp_name */
1754 sizeof(repeatobject), /* tp_basicsize */
1755 0, /* tp_itemsize */
1756 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001757 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758 0, /* tp_print */
1759 0, /* tp_getattr */
1760 0, /* tp_setattr */
1761 0, /* tp_compare */
1762 0, /* tp_repr */
1763 0, /* tp_as_number */
1764 0, /* tp_as_sequence */
1765 0, /* tp_as_mapping */
1766 0, /* tp_hash */
1767 0, /* tp_call */
1768 0, /* tp_str */
1769 PyObject_GenericGetAttr, /* tp_getattro */
1770 0, /* tp_setattro */
1771 0, /* tp_as_buffer */
1772 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1773 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001774 repeat_doc, /* tp_doc */
1775 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001776 0, /* tp_clear */
1777 0, /* tp_richcompare */
1778 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001779 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001780 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001781 0, /* tp_methods */
1782 0, /* tp_members */
1783 0, /* tp_getset */
1784 0, /* tp_base */
1785 0, /* tp_dict */
1786 0, /* tp_descr_get */
1787 0, /* tp_descr_set */
1788 0, /* tp_dictoffset */
1789 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001790 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001791 repeat_new, /* tp_new */
1792 PyObject_GC_Del, /* tp_free */
1793};
1794
1795
1796/* module level code ********************************************************/
1797
1798PyDoc_STRVAR(module_doc,
1799"Functional tools for creating and using iterators.\n\
1800\n\
1801Infinite iterators:\n\
1802count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001803cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00001804repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001805\n\
1806Iterators terminating on the shortest input sequence:\n\
1807izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001808ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
1809ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001810islice(seq, [start,] stop [, step]) --> elements from\n\
1811 seq[start:stop:step]\n\
1812imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
1813starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001814chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001815takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
1816dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
1817");
1818
1819
1820PyMODINIT_FUNC
1821inititertools(void)
1822{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001823 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001824 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001825 char *name;
1826 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001827 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001828 &dropwhile_type,
1829 &takewhile_type,
1830 &islice_type,
1831 &starmap_type,
1832 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001833 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001834 &ifilter_type,
1835 &ifilterfalse_type,
1836 &count_type,
1837 &izip_type,
1838 &repeat_type,
1839 NULL
1840 };
1841
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001842 m = Py_InitModule3("itertools", NULL, module_doc);
1843
Raymond Hettinger60eca932003-02-09 06:40:58 +00001844 for (i=0 ; typelist[i] != NULL ; i++) {
1845 if (PyType_Ready(typelist[i]) < 0)
1846 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00001847 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001848 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001849 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00001850 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001851 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001852}