?  search:
summary | log | graphiclog1 | graphiclog2 | commit | commitdiff | tree | refs | edit | fork
blame | history | raw | HEAD
Minor typo fixes
[pytest.git] / Modules / itertoolsmodule.c
blobd91389033ff4ec719105716ef4dafc8b0dbbd16d
1
2 #include "Python.h"
3 #include "structmember.h"
4
5 /* Itertools module written and maintained
6 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9 */
10
11
12 /* groupby object ***********************************************************/
13
14 typedef struct {
15 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
21 } groupbyobject;
22
23 static PyTypeObject groupby_type;
24 static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26 static PyObject *
27 groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28 {
29 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
32
33 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
51 }
52
53 static void
54 groupby_dealloc(groupbyobject *gbo)
55 {
56 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 gbo->ob_type->tp_free(gbo);
63 }
64
65 static int
66 groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67 {
68 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
74 }
75
76 static PyObject *
77 groupby_next(groupbyobject *gbo)
78 {
79 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
80
81 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
89
90 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
97
98 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
101
102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
113
114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
117
118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
122
123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
127
128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
131
132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
135 }
136
137 PyDoc_STRVAR(groupby_doc,
138 "groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139 (key, sub-iterator) grouped by each value of key(value).\n");
140
141 static PyTypeObject groupby_type = {
142 PyObject_HEAD_INIT(NULL)
143 0, /* ob_size */
144 "itertools.groupby", /* tp_name */
145 sizeof(groupbyobject), /* tp_basicsize */
146 0, /* tp_itemsize */
147 /* methods */
148 (destructor)groupby_dealloc, /* tp_dealloc */
149 0, /* tp_print */
150 0, /* tp_getattr */
151 0, /* tp_setattr */
152 0, /* tp_compare */
153 0, /* tp_repr */
154 0, /* tp_as_number */
155 0, /* tp_as_sequence */
156 0, /* tp_as_mapping */
157 0, /* tp_hash */
158 0, /* tp_call */
159 0, /* tp_str */
160 PyObject_GenericGetAttr, /* tp_getattro */
161 0, /* tp_setattro */
162 0, /* tp_as_buffer */
163 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
164 Py_TPFLAGS_BASETYPE, /* tp_flags */
165 groupby_doc, /* tp_doc */
166 (traverseproc)groupby_traverse, /* tp_traverse */
167 0, /* tp_clear */
168 0, /* tp_richcompare */
169 0, /* tp_weaklistoffset */
170 PyObject_SelfIter, /* tp_iter */
171 (iternextfunc)groupby_next, /* tp_iternext */
172 0, /* tp_methods */
173 0, /* tp_members */
174 0, /* tp_getset */
175 0, /* tp_base */
176 0, /* tp_dict */
177 0, /* tp_descr_get */
178 0, /* tp_descr_set */
179 0, /* tp_dictoffset */
180 0, /* tp_init */
181 0, /* tp_alloc */
182 groupby_new, /* tp_new */
183 PyObject_GC_Del, /* tp_free */
184 };
185
186
187 /* _grouper object (internal) ************************************************/
188
189 typedef struct {
190 PyObject_HEAD
191 PyObject *parent;
192 PyObject *tgtkey;
193 } _grouperobject;
194
195 static PyTypeObject _grouper_type;
196
197 static PyObject *
198 _grouper_create(groupbyobject *parent, PyObject *tgtkey)
199 {
200 _grouperobject *igo;
201
202 igo = PyObject_New(_grouperobject, &_grouper_type);
203 if (igo == NULL)
204 return NULL;
205 igo->parent = (PyObject *)parent;
206 Py_INCREF(parent);
207 igo->tgtkey = tgtkey;
208 Py_INCREF(tgtkey);
209
210 return (PyObject *)igo;
211 }
212
213 static void
214 _grouper_dealloc(_grouperobject *igo)
215 {
216 Py_DECREF(igo->parent);
217 Py_DECREF(igo->tgtkey);
218 PyObject_Del(igo);
219 }
220
221 static PyObject *
222 _grouper_next(_grouperobject *igo)
223 {
224 groupbyobject *gbo = (groupbyobject *)igo->parent;
225 PyObject *newvalue, *newkey, *r;
226 int rcmp;
227
228 if (gbo->currvalue == NULL) {
229 newvalue = PyIter_Next(gbo->it);
230 if (newvalue == NULL)
231 return NULL;
232
233 if (gbo->keyfunc == Py_None) {
234 newkey = newvalue;
235 Py_INCREF(newvalue);
236 } else {
237 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
238 newvalue, NULL);
239 if (newkey == NULL) {
240 Py_DECREF(newvalue);
241 return NULL;
242 }
243 }
244
245 assert(gbo->currkey == NULL);
246 gbo->currkey = newkey;
247 gbo->currvalue = newvalue;
248 }
249
250 assert(gbo->currkey != NULL);
251 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
252 if (rcmp <= 0)
253 /* got any error or current group is end */
254 return NULL;
255
256 r = gbo->currvalue;
257 gbo->currvalue = NULL;
258 Py_CLEAR(gbo->currkey);
259
260 return r;
261 }
262
263 static PyTypeObject _grouper_type = {
264 PyObject_HEAD_INIT(NULL)
265 0, /* ob_size */
266 "itertools._grouper", /* tp_name */
267 sizeof(_grouperobject), /* tp_basicsize */
268 0, /* tp_itemsize */
269 /* methods */
270 (destructor)_grouper_dealloc, /* tp_dealloc */
271 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, /* tp_flags */
286 0, /* tp_doc */
287 0, /* tp_traverse */
288 0, /* tp_clear */
289 0, /* tp_richcompare */
290 0, /* tp_weaklistoffset */
291 PyObject_SelfIter, /* tp_iter */
292 (iternextfunc)_grouper_next, /* tp_iternext */
293 0, /* tp_methods */
294 0, /* tp_members */
295 0, /* tp_getset */
296 0, /* tp_base */
297 0, /* tp_dict */
298 0, /* tp_descr_get */
299 0, /* tp_descr_set */
300 0, /* tp_dictoffset */
301 0, /* tp_init */
302 0, /* tp_alloc */
303 0, /* tp_new */
304 PyObject_Del, /* tp_free */
305 };
306
307
308
309 /* tee object and with supporting function and objects ***************/
310
311 /* The teedataobject pre-allocates space for LINKCELLS number of objects.
312 To help the object fit neatly inside cache lines (space for 16 to 32
313 pointers), the value should be a multiple of 16 minus space for
314 the other structure members including PyHEAD overhead. The larger the
315 value, the less memory overhead per object and the less time spent
316 allocating/deallocating new links. The smaller the number, the less
317 wasted space and the more rapid freeing of older data.
318 */
319 #define LINKCELLS 57
320
321 typedef struct {
322 PyObject_HEAD
323 PyObject *it;
324 int numread;
325 PyObject *nextlink;
326 PyObject *(values[LINKCELLS]);
327 } teedataobject;
328
329 typedef struct {
330 PyObject_HEAD
331 teedataobject *dataobj;
332 int index;
333 PyObject *weakreflist;
334 } teeobject;
335
336 static PyTypeObject teedataobject_type;
337
338 static PyObject *
339 teedataobject_new(PyObject *it)
340 {
341 teedataobject *tdo;
342
343 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
344 if (tdo == NULL)
345 return NULL;
346
347 tdo->numread = 0;
348 tdo->nextlink = NULL;
349 Py_INCREF(it);
350 tdo->it = it;
351 PyObject_GC_Track(tdo);
352 return (PyObject *)tdo;
353 }
354
355 static PyObject *
356 teedataobject_jumplink(teedataobject *tdo)
357 {
358 if (tdo->nextlink == NULL)
359 tdo->nextlink = teedataobject_new(tdo->it);
360 Py_XINCREF(tdo->nextlink);
361 return tdo->nextlink;
362 }
363
364 static PyObject *
365 teedataobject_getitem(teedataobject *tdo, int i)
366 {
367 PyObject *value;
368
369 assert(i < LINKCELLS);
370 if (i < tdo->numread)
371 value = tdo->values[i];
372 else {
373 /* this is the lead iterator, so fetch more data */
374 assert(i == tdo->numread);
375 value = PyIter_Next(tdo->it);
376 if (value == NULL)
377 return NULL;
378 tdo->numread++;
379 tdo->values[i] = value;
380 }
381 Py_INCREF(value);
382 return value;
383 }
384
385 static int
386 teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
387 {
388 int i;
389 Py_VISIT(tdo->it);
390 for (i = 0; i < tdo->numread; i++)
391 Py_VISIT(tdo->values[i]);
392 Py_VISIT(tdo->nextlink);
393 return 0;
394 }
395
396 static int
397 teedataobject_clear(teedataobject *tdo)
398 {
399 int i;
400 Py_CLEAR(tdo->it);
401 for (i=0 ; i<tdo->numread ; i++)
402 Py_CLEAR(tdo->values[i]);
403 Py_CLEAR(tdo->nextlink);
404 return 0;
405 }
406
407 static void
408 teedataobject_dealloc(teedataobject *tdo)
409 {
410 PyObject_GC_UnTrack(tdo);
411 teedataobject_clear(tdo);
412 PyObject_GC_Del(tdo);
413 }
414
415 PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
416
417 static PyTypeObject teedataobject_type = {
418 PyObject_HEAD_INIT(0) /* Must fill in type value later */
419 0, /* ob_size */
420 "itertools.tee_dataobject", /* tp_name */
421 sizeof(teedataobject), /* tp_basicsize */
422 0, /* tp_itemsize */
423 /* methods */
424 (destructor)teedataobject_dealloc, /* tp_dealloc */
425 0, /* tp_print */
426 0, /* tp_getattr */
427 0, /* tp_setattr */
428 0, /* tp_compare */
429 0, /* tp_repr */
430 0, /* tp_as_number */
431 0, /* tp_as_sequence */
432 0, /* tp_as_mapping */
433 0, /* tp_hash */
434 0, /* tp_call */
435 0, /* tp_str */
436 PyObject_GenericGetAttr, /* tp_getattro */
437 0, /* tp_setattro */
438 0, /* tp_as_buffer */
439 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
440 teedataobject_doc, /* tp_doc */
441 (traverseproc)teedataobject_traverse, /* tp_traverse */
442 (inquiry)teedataobject_clear, /* tp_clear */
443 0, /* tp_richcompare */
444 0, /* tp_weaklistoffset */
445 0, /* tp_iter */
446 0, /* tp_iternext */
447 0, /* tp_methods */
448 0, /* tp_members */
449 0, /* tp_getset */
450 0, /* tp_base */
451 0, /* tp_dict */
452 0, /* tp_descr_get */
453 0, /* tp_descr_set */
454 0, /* tp_dictoffset */
455 0, /* tp_init */
456 0, /* tp_alloc */
457 0, /* tp_new */
458 PyObject_GC_Del, /* tp_free */
459 };
460
461
462 static PyTypeObject tee_type;
463
464 static PyObject *
465 tee_next(teeobject *to)
466 {
467 PyObject *value, *link;
468
469 if (to->index >= LINKCELLS) {
470 link = teedataobject_jumplink(to->dataobj);
471 Py_DECREF(to->dataobj);
472 to->dataobj = (teedataobject *)link;
473 to->index = 0;
474 }
475 value = teedataobject_getitem(to->dataobj, to->index);
476 if (value == NULL)
477 return NULL;
478 to->index++;
479 return value;
480 }
481
482 static int
483 tee_traverse(teeobject *to, visitproc visit, void *arg)
484 {
485 Py_VISIT((PyObject *)to->dataobj);
486 return 0;
487 }
488
489 static PyObject *
490 tee_copy(teeobject *to)
491 {
492 teeobject *newto;
493
494 newto = PyObject_GC_New(teeobject, &tee_type);
495 if (newto == NULL)
496 return NULL;
497 Py_INCREF(to->dataobj);
498 newto->dataobj = to->dataobj;
499 newto->index = to->index;
500 newto->weakreflist = NULL;
501 PyObject_GC_Track(newto);
502 return (PyObject *)newto;
503 }
504
505 PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
506
507 static PyObject *
508 tee_fromiterable(PyObject *iterable)
509 {
510 teeobject *to;
511 PyObject *it = NULL;
512
513 it = PyObject_GetIter(iterable);
514 if (it == NULL)
515 return NULL;
516 if (PyObject_TypeCheck(it, &tee_type)) {
517 to = (teeobject *)tee_copy((teeobject *)it);
518 goto done;
519 }
520
521 to = PyObject_GC_New(teeobject, &tee_type);
522 if (to == NULL)
523 goto done;
524 to->dataobj = (teedataobject *)teedataobject_new(it);
525 if (!to->dataobj) {
526 PyObject_GC_Del(to);
527 to = NULL;
528 goto done;
529 }
530
531 to->index = 0;
532 to->weakreflist = NULL;
533 PyObject_GC_Track(to);
534 done:
535 Py_XDECREF(it);
536 return (PyObject *)to;
537 }
538
539 static PyObject *
540 tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
541 {
542 PyObject *iterable;
543
544 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
545 return NULL;
546 return tee_fromiterable(iterable);
547 }
548
549 static int
550 tee_clear(teeobject *to)
551 {
552 if (to->weakreflist != NULL)
553 PyObject_ClearWeakRefs((PyObject *) to);
554 Py_CLEAR(to->dataobj);
555 return 0;
556 }
557
558 static void
559 tee_dealloc(teeobject *to)
560 {
561 PyObject_GC_UnTrack(to);
562 tee_clear(to);
563 PyObject_GC_Del(to);
564 }
565
566 PyDoc_STRVAR(teeobject_doc,
567 "Iterator wrapped to make it copyable");
568
569 static PyMethodDef tee_methods[] = {
570 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
571 {NULL, NULL} /* sentinel */
572 };
573
574 static PyTypeObject tee_type = {
575 PyObject_HEAD_INIT(NULL)
576 0, /* ob_size */
577 "itertools.tee", /* tp_name */
578 sizeof(teeobject), /* tp_basicsize */
579 0, /* tp_itemsize */
580 /* methods */
581 (destructor)tee_dealloc, /* tp_dealloc */
582 0, /* tp_print */
583 0, /* tp_getattr */
584 0, /* tp_setattr */
585 0, /* tp_compare */
586 0, /* tp_repr */
587 0, /* tp_as_number */
588 0, /* tp_as_sequence */
589 0, /* tp_as_mapping */
590 0, /* tp_hash */
591 0, /* tp_call */
592 0, /* tp_str */
593 0, /* tp_getattro */
594 0, /* tp_setattro */
595 0, /* tp_as_buffer */
596 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
597 teeobject_doc, /* tp_doc */
598 (traverseproc)tee_traverse, /* tp_traverse */
599 (inquiry)tee_clear, /* tp_clear */
600 0, /* tp_richcompare */
601 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
602 PyObject_SelfIter, /* tp_iter */
603 (iternextfunc)tee_next, /* tp_iternext */
604 tee_methods, /* tp_methods */
605 0, /* tp_members */
606 0, /* tp_getset */
607 0, /* tp_base */
608 0, /* tp_dict */
609 0, /* tp_descr_get */
610 0, /* tp_descr_set */
611 0, /* tp_dictoffset */
612 0, /* tp_init */
613 0, /* tp_alloc */
614 tee_new, /* tp_new */
615 PyObject_GC_Del, /* tp_free */
616 };
617
618 static PyObject *
619 tee(PyObject *self, PyObject *args)
620 {
621 int i, n=2;
622 PyObject *it, *iterable, *copyable, *result;
623
624 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
625 return NULL;
626 result = PyTuple_New(n);
627 if (result == NULL)
628 return NULL;
629 if (n == 0)
630 return result;
631 it = PyObject_GetIter(iterable);
632 if (it == NULL) {
633 Py_DECREF(result);
634 return NULL;
635 }
636 if (!PyObject_HasAttrString(it, "__copy__")) {
637 copyable = tee_fromiterable(it);
638 Py_DECREF(it);
639 if (copyable == NULL) {
640 Py_DECREF(result);
641 return NULL;
642 }
643 } else
644 copyable = it;
645 PyTuple_SET_ITEM(result, 0, copyable);
646 for (i=1 ; i<n ; i++) {
647 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
648 if (copyable == NULL) {
649 Py_DECREF(result);
650 return NULL;
651 }
652 PyTuple_SET_ITEM(result, i, copyable);
653 }
654 return result;
655 }
656
657 PyDoc_STRVAR(tee_doc,
658 "tee(iterable, n=2) --> tuple of n independent iterators.");
659
660
661 /* cycle object **********************************************************/
662
663 typedef struct {
664 PyObject_HEAD
665 PyObject *it;
666 PyObject *saved;
667 int firstpass;
668 } cycleobject;
669
670 static PyTypeObject cycle_type;
671
672 static PyObject *
673 cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
674 {
675 PyObject *it;
676 PyObject *iterable;
677 PyObject *saved;
678 cycleobject *lz;
679
680 if (!_PyArg_NoKeywords("cycle()", kwds))
681 return NULL;
682
683 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
684 return NULL;
685
686 /* Get iterator. */
687 it = PyObject_GetIter(iterable);
688 if (it == NULL)
689 return NULL;
690
691 saved = PyList_New(0);
692 if (saved == NULL) {
693 Py_DECREF(it);
694 return NULL;
695 }
696
697 /* create cycleobject structure */
698 lz = (cycleobject *)type->tp_alloc(type, 0);
699 if (lz == NULL) {
700 Py_DECREF(it);
701 Py_DECREF(saved);
702 return NULL;
703 }
704 lz->it = it;
705 lz->saved = saved;
706 lz->firstpass = 0;
707
708 return (PyObject *)lz;
709 }
710
711 static void
712 cycle_dealloc(cycleobject *lz)
713 {
714 PyObject_GC_UnTrack(lz);
715 Py_XDECREF(lz->saved);
716 Py_XDECREF(lz->it);
717 lz->ob_type->tp_free(lz);
718 }
719
720 static int
721 cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
722 {
723 Py_VISIT(lz->it);
724 Py_VISIT(lz->saved);
725 return 0;
726 }
727
728 static PyObject *
729 cycle_next(cycleobject *lz)
730 {
731 PyObject *item;
732 PyObject *it;
733 PyObject *tmp;
734
735 while (1) {
736 item = PyIter_Next(lz->it);
737 if (item != NULL) {
738 if (!lz->firstpass)
739 PyList_Append(lz->saved, item);
740 return item;
741 }
742 if (PyErr_Occurred()) {
743 if (PyErr_ExceptionMatches(PyExc_StopIteration))
744 PyErr_Clear();
745 else
746 return NULL;
747 }
748 if (PyList_Size(lz->saved) == 0)
749 return NULL;
750 it = PyObject_GetIter(lz->saved);
751 if (it == NULL)
752 return NULL;
753 tmp = lz->it;
754 lz->it = it;
755 lz->firstpass = 1;
756 Py_DECREF(tmp);
757 }
758 }
759
760 PyDoc_STRVAR(cycle_doc,
761 "cycle(iterable) --> cycle object\n\
762 \n\
763 Return elements from the iterable until it is exhausted.\n\
764 Then repeat the sequence indefinitely.");
765
766 static PyTypeObject cycle_type = {
767 PyObject_HEAD_INIT(NULL)
768 0, /* ob_size */
769 "itertools.cycle", /* tp_name */
770 sizeof(cycleobject), /* tp_basicsize */
771 0, /* tp_itemsize */
772 /* methods */
773 (destructor)cycle_dealloc, /* tp_dealloc */
774 0, /* tp_print */
775 0, /* tp_getattr */
776 0, /* tp_setattr */
777 0, /* tp_compare */
778 0, /* tp_repr */
779 0, /* tp_as_number */
780 0, /* tp_as_sequence */
781 0, /* tp_as_mapping */
782 0, /* tp_hash */
783 0, /* tp_call */
784 0, /* tp_str */
785 PyObject_GenericGetAttr, /* tp_getattro */
786 0, /* tp_setattro */
787 0, /* tp_as_buffer */
788 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
789 Py_TPFLAGS_BASETYPE, /* tp_flags */
790 cycle_doc, /* tp_doc */
791 (traverseproc)cycle_traverse, /* tp_traverse */
792 0, /* tp_clear */
793 0, /* tp_richcompare */
794 0, /* tp_weaklistoffset */
795 PyObject_SelfIter, /* tp_iter */
796 (iternextfunc)cycle_next, /* tp_iternext */
797 0, /* tp_methods */
798 0, /* tp_members */
799 0, /* tp_getset */
800 0, /* tp_base */
801 0, /* tp_dict */
802 0, /* tp_descr_get */
803 0, /* tp_descr_set */
804 0, /* tp_dictoffset */
805 0, /* tp_init */
806 0, /* tp_alloc */
807 cycle_new, /* tp_new */
808 PyObject_GC_Del, /* tp_free */
809 };
810
811
812 /* dropwhile object **********************************************************/
813
814 typedef struct {
815 PyObject_HEAD
816 PyObject *func;
817 PyObject *it;
818 long start;
819 } dropwhileobject;
820
821 static PyTypeObject dropwhile_type;
822
823 static PyObject *
824 dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
825 {
826 PyObject *func, *seq;
827 PyObject *it;
828 dropwhileobject *lz;
829
830 if (!_PyArg_NoKeywords("dropwhile()", kwds))
831 return NULL;
832
833 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
834 return NULL;
835
836 /* Get iterator. */
837 it = PyObject_GetIter(seq);
838 if (it == NULL)
839 return NULL;
840
841 /* create dropwhileobject structure */
842 lz = (dropwhileobject *)type->tp_alloc(type, 0);
843 if (lz == NULL) {
844 Py_DECREF(it);
845 return NULL;
846 }
847 Py_INCREF(func);
848 lz->func = func;
849 lz->it = it;
850 lz->start = 0;
851
852 return (PyObject *)lz;
853 }
854
855 static void
856 dropwhile_dealloc(dropwhileobject *lz)
857 {
858 PyObject_GC_UnTrack(lz);
859 Py_XDECREF(lz->func);
860 Py_XDECREF(lz->it);
861 lz->ob_type->tp_free(lz);
862 }
863
864 static int
865 dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
866 {
867 Py_VISIT(lz->it);
868 Py_VISIT(lz->func);
869 return 0;
870 }
871
872 static PyObject *
873 dropwhile_next(dropwhileobject *lz)
874 {
875 PyObject *item, *good;
876 PyObject *it = lz->it;
877 long ok;
878 PyObject *(*iternext)(PyObject *);
879
880 assert(PyIter_Check(it));
881 iternext = *it->ob_type->tp_iternext;
882 for (;;) {
883 item = iternext(it);
884 if (item == NULL)
885 return NULL;
886 if (lz->start == 1)
887 return item;
888
889 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
890 if (good == NULL) {
891 Py_DECREF(item);
892 return NULL;
893 }
894 ok = PyObject_IsTrue(good);
895 Py_DECREF(good);
896 if (!ok) {
897 lz->start = 1;
898 return item;
899 }
900 Py_DECREF(item);
901 }
902 }
903
904 PyDoc_STRVAR(dropwhile_doc,
905 "dropwhile(predicate, iterable) --> dropwhile object\n\
906 \n\
907 Drop items from the iterable while predicate(item) is true.\n\
908 Afterwards, return every element until the iterable is exhausted.");
909
910 static PyTypeObject dropwhile_type = {
911 PyObject_HEAD_INIT(NULL)
912 0, /* ob_size */
913 "itertools.dropwhile", /* tp_name */
914 sizeof(dropwhileobject), /* tp_basicsize */
915 0, /* tp_itemsize */
916 /* methods */
917 (destructor)dropwhile_dealloc, /* tp_dealloc */
918 0, /* tp_print */
919 0, /* tp_getattr */
920 0, /* tp_setattr */
921 0, /* tp_compare */
922 0, /* tp_repr */
923 0, /* tp_as_number */
924 0, /* tp_as_sequence */
925 0, /* tp_as_mapping */
926 0, /* tp_hash */
927 0, /* tp_call */
928 0, /* tp_str */
929 PyObject_GenericGetAttr, /* tp_getattro */
930 0, /* tp_setattro */
931 0, /* tp_as_buffer */
932 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
933 Py_TPFLAGS_BASETYPE, /* tp_flags */
934 dropwhile_doc, /* tp_doc */
935 (traverseproc)dropwhile_traverse, /* tp_traverse */
936 0, /* tp_clear */
937 0, /* tp_richcompare */
938 0, /* tp_weaklistoffset */
939 PyObject_SelfIter, /* tp_iter */
940 (iternextfunc)dropwhile_next, /* tp_iternext */
941 0, /* tp_methods */
942 0, /* tp_members */
943 0, /* tp_getset */
944 0, /* tp_base */
945 0, /* tp_dict */
946 0, /* tp_descr_get */
947 0, /* tp_descr_set */
948 0, /* tp_dictoffset */
949 0, /* tp_init */
950 0, /* tp_alloc */
951 dropwhile_new, /* tp_new */
952 PyObject_GC_Del, /* tp_free */
953 };
954
955
956 /* takewhile object **********************************************************/
957
958 typedef struct {
959 PyObject_HEAD
960 PyObject *func;
961 PyObject *it;
962 long stop;
963 } takewhileobject;
964
965 static PyTypeObject takewhile_type;
966
967 static PyObject *
968 takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
969 {
970 PyObject *func, *seq;
971 PyObject *it;
972 takewhileobject *lz;
973
974 if (!_PyArg_NoKeywords("takewhile()", kwds))
975 return NULL;
976
977 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
978 return NULL;
979
980 /* Get iterator. */
981 it = PyObject_GetIter(seq);
982 if (it == NULL)
983 return NULL;
984
985 /* create takewhileobject structure */
986 lz = (takewhileobject *)type->tp_alloc(type, 0);
987 if (lz == NULL) {
988 Py_DECREF(it);
989 return NULL;
990 }
991 Py_INCREF(func);
992 lz->func = func;
993 lz->it = it;
994 lz->stop = 0;
995
996 return (PyObject *)lz;
997 }
998
999 static void
1000 takewhile_dealloc(takewhileobject *lz)
1001 {
1002 PyObject_GC_UnTrack(lz);
1003 Py_XDECREF(lz->func);
1004 Py_XDECREF(lz->it);
1005 lz->ob_type->tp_free(lz);
1006 }
1007
1008 static int
1009 takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1010 {
1011 Py_VISIT(lz->it);
1012 Py_VISIT(lz->func);
1013 return 0;
1014 }
1015
1016 static PyObject *
1017 takewhile_next(takewhileobject *lz)
1018 {
1019 PyObject *item, *good;
1020 PyObject *it = lz->it;
1021 long ok;
1022
1023 if (lz->stop == 1)
1024 return NULL;
1025
1026 assert(PyIter_Check(it));
1027 item = (*it->ob_type->tp_iternext)(it);
1028 if (item == NULL)
1029 return NULL;
1030
1031 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1032 if (good == NULL) {
1033 Py_DECREF(item);
1034 return NULL;
1035 }
1036 ok = PyObject_IsTrue(good);
1037 Py_DECREF(good);
1038 if (ok)
1039 return item;
1040 Py_DECREF(item);
1041 lz->stop = 1;
1042 return NULL;
1043 }
1044
1045 PyDoc_STRVAR(takewhile_doc,
1046 "takewhile(predicate, iterable) --> takewhile object\n\
1047 \n\
1048 Return successive entries from an iterable as long as the \n\
1049 predicate evaluates to true for each entry.");
1050
1051 static PyTypeObject takewhile_type = {
1052 PyObject_HEAD_INIT(NULL)
1053 0, /* ob_size */
1054 "itertools.takewhile", /* tp_name */
1055 sizeof(takewhileobject), /* tp_basicsize */
1056 0, /* tp_itemsize */
1057 /* methods */
1058 (destructor)takewhile_dealloc, /* tp_dealloc */
1059 0, /* tp_print */
1060 0, /* tp_getattr */
1061 0, /* tp_setattr */
1062 0, /* tp_compare */
1063 0, /* tp_repr */
1064 0, /* tp_as_number */
1065 0, /* tp_as_sequence */
1066 0, /* tp_as_mapping */
1067 0, /* tp_hash */
1068 0, /* tp_call */
1069 0, /* tp_str */
1070 PyObject_GenericGetAttr, /* tp_getattro */
1071 0, /* tp_setattro */
1072 0, /* tp_as_buffer */
1073 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1074 Py_TPFLAGS_BASETYPE, /* tp_flags */
1075 takewhile_doc, /* tp_doc */
1076 (traverseproc)takewhile_traverse, /* tp_traverse */
1077 0, /* tp_clear */
1078 0, /* tp_richcompare */
1079 0, /* tp_weaklistoffset */
1080 PyObject_SelfIter, /* tp_iter */
1081 (iternextfunc)takewhile_next, /* tp_iternext */
1082 0, /* tp_methods */
1083 0, /* tp_members */
1084 0, /* tp_getset */
1085 0, /* tp_base */
1086 0, /* tp_dict */
1087 0, /* tp_descr_get */
1088 0, /* tp_descr_set */
1089 0, /* tp_dictoffset */
1090 0, /* tp_init */
1091 0, /* tp_alloc */
1092 takewhile_new, /* tp_new */
1093 PyObject_GC_Del, /* tp_free */
1094 };
1095
1096
1097 /* islice object ************************************************************/
1098
1099 typedef struct {
1100 PyObject_HEAD
1101 PyObject *it;
1102 Py_ssize_t next;
1103 Py_ssize_t stop;
1104 Py_ssize_t step;
1105 Py_ssize_t cnt;
1106 } isliceobject;
1107
1108 static PyTypeObject islice_type;
1109
1110 static PyObject *
1111 islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1112 {
1113 PyObject *seq;
1114 Py_ssize_t start=0, stop=-1, step=1;
1115 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1116 Py_ssize_t numargs;
1117 isliceobject *lz;
1118
1119 if (!_PyArg_NoKeywords("islice()", kwds))
1120 return NULL;
1121
1122 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1123 return NULL;
1124
1125 numargs = PyTuple_Size(args);
1126 if (numargs == 2) {
1127 if (a1 != Py_None) {
1128 stop = PyInt_AsSsize_t(a1);
1129 if (stop == -1) {
1130 if (PyErr_Occurred())
1131 PyErr_Clear();
1132 PyErr_SetString(PyExc_ValueError,
1133 "Stop argument for islice() must be a non-negative integer or None.");
1134 return NULL;
1135 }
1136 }
1137 } else {
1138 if (a1 != Py_None)
1139 start = PyInt_AsSsize_t(a1);
1140 if (start == -1 && PyErr_Occurred())
1141 PyErr_Clear();
1142 if (a2 != Py_None) {
1143 stop = PyInt_AsSsize_t(a2);
1144 if (stop == -1) {
1145 if (PyErr_Occurred())
1146 PyErr_Clear();
1147 PyErr_SetString(PyExc_ValueError,
1148 "Stop argument for islice() must be a non-negative integer or None.");
1149 return NULL;
1150 }
1151 }
1152 }
1153 if (start<0 || stop<-1) {
1154 PyErr_SetString(PyExc_ValueError,
1155 "Indices for islice() must be non-negative integers or None.");
1156 return NULL;
1157 }
1158
1159 if (a3 != NULL) {
1160 if (a3 != Py_None)
1161 step = PyInt_AsSsize_t(a3);
1162 if (step == -1 && PyErr_Occurred())
1163 PyErr_Clear();
1164 }
1165 if (step<1) {
1166 PyErr_SetString(PyExc_ValueError,
1167 "Step for islice() must be a positive integer or None.");
1168 return NULL;
1169 }
1170
1171 /* Get iterator. */
1172 it = PyObject_GetIter(seq);
1173 if (it == NULL)
1174 return NULL;
1175
1176 /* create isliceobject structure */
1177 lz = (isliceobject *)type->tp_alloc(type, 0);
1178 if (lz == NULL) {
1179 Py_DECREF(it);
1180 return NULL;
1181 }
1182 lz->it = it;
1183 lz->next = start;
1184 lz->stop = stop;
1185 lz->step = step;
1186 lz->cnt = 0L;
1187
1188 return (PyObject *)lz;
1189 }
1190
1191 static void
1192 islice_dealloc(isliceobject *lz)
1193 {
1194 PyObject_GC_UnTrack(lz);
1195 Py_XDECREF(lz->it);
1196 lz->ob_type->tp_free(lz);
1197 }
1198
1199 static int
1200 islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1201 {
1202 Py_VISIT(lz->it);
1203 return 0;
1204 }
1205
1206 static PyObject *
1207 islice_next(isliceobject *lz)
1208 {
1209 PyObject *item;
1210 PyObject *it = lz->it;
1211 Py_ssize_t oldnext;
1212 PyObject *(*iternext)(PyObject *);
1213
1214 assert(PyIter_Check(it));
1215 iternext = *it->ob_type->tp_iternext;
1216 while (lz->cnt < lz->next) {
1217 item = iternext(it);
1218 if (item == NULL)
1219 return NULL;
1220 Py_DECREF(item);
1221 lz->cnt++;
1222 }
1223 if (lz->stop != -1 && lz->cnt >= lz->stop)
1224 return NULL;
1225 assert(PyIter_Check(it));
1226 item = iternext(it);
1227 if (item == NULL)
1228 return NULL;
1229 lz->cnt++;
1230 oldnext = lz->next;
1231 lz->next += lz->step;
1232 if (lz->next < oldnext) /* Check for overflow */
1233 lz->next = lz->stop;
1234 return item;
1235 }
1236
1237 PyDoc_STRVAR(islice_doc,
1238 "islice(iterable, [start,] stop [, step]) --> islice object\n\
1239 \n\
1240 Return an iterator whose next() method returns selected values from an\n\
1241 iterable. If start is specified, will skip all preceding elements;\n\
1242 otherwise, start defaults to zero. Step defaults to one. If\n\
1243 specified as another value, step determines how many values are \n\
1244 skipped between successive calls. Works like a slice() on a list\n\
1245 but returns an iterator.");
1246
1247 static PyTypeObject islice_type = {
1248 PyObject_HEAD_INIT(NULL)
1249 0, /* ob_size */
1250 "itertools.islice", /* tp_name */
1251 sizeof(isliceobject), /* tp_basicsize */
1252 0, /* tp_itemsize */
1253 /* methods */
1254 (destructor)islice_dealloc, /* tp_dealloc */
1255 0, /* tp_print */
1256 0, /* tp_getattr */
1257 0, /* tp_setattr */
1258 0, /* tp_compare */
1259 0, /* tp_repr */
1260 0, /* tp_as_number */
1261 0, /* tp_as_sequence */
1262 0, /* tp_as_mapping */
1263 0, /* tp_hash */
1264 0, /* tp_call */
1265 0, /* tp_str */
1266 PyObject_GenericGetAttr, /* tp_getattro */
1267 0, /* tp_setattro */
1268 0, /* tp_as_buffer */
1269 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1270 Py_TPFLAGS_BASETYPE, /* tp_flags */
1271 islice_doc, /* tp_doc */
1272 (traverseproc)islice_traverse, /* tp_traverse */
1273 0, /* tp_clear */
1274 0, /* tp_richcompare */
1275 0, /* tp_weaklistoffset */
1276 PyObject_SelfIter, /* tp_iter */
1277 (iternextfunc)islice_next, /* tp_iternext */
1278 0, /* tp_methods */
1279 0, /* tp_members */
1280 0, /* tp_getset */
1281 0, /* tp_base */
1282 0, /* tp_dict */
1283 0, /* tp_descr_get */
1284 0, /* tp_descr_set */
1285 0, /* tp_dictoffset */
1286 0, /* tp_init */
1287 0, /* tp_alloc */
1288 islice_new, /* tp_new */
1289 PyObject_GC_Del, /* tp_free */
1290 };
1291
1292
1293 /* starmap object ************************************************************/
1294
1295 typedef struct {
1296 PyObject_HEAD
1297 PyObject *func;
1298 PyObject *it;
1299 } starmapobject;
1300
1301 static PyTypeObject starmap_type;
1302
1303 static PyObject *
1304 starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1305 {
1306 PyObject *func, *seq;
1307 PyObject *it;
1308 starmapobject *lz;
1309
1310 if (!_PyArg_NoKeywords("starmap()", kwds))
1311 return NULL;
1312
1313 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1314 return NULL;
1315
1316 /* Get iterator. */
1317 it = PyObject_GetIter(seq);
1318 if (it == NULL)
1319 return NULL;
1320
1321 /* create starmapobject structure */
1322 lz = (starmapobject *)type->tp_alloc(type, 0);
1323 if (lz == NULL) {
1324 Py_DECREF(it);
1325 return NULL;
1326 }
1327 Py_INCREF(func);
1328 lz->func = func;
1329 lz->it = it;
1330
1331 return (PyObject *)lz;
1332 }
1333
1334 static void
1335 starmap_dealloc(starmapobject *lz)
1336 {
1337 PyObject_GC_UnTrack(lz);
1338 Py_XDECREF(lz->func);
1339 Py_XDECREF(lz->it);
1340 lz->ob_type->tp_free(lz);
1341 }
1342
1343 static int
1344 starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1345 {
1346 Py_VISIT(lz->it);
1347 Py_VISIT(lz->func);
1348 return 0;
1349 }
1350
1351 static PyObject *
1352 starmap_next(starmapobject *lz)
1353 {
1354 PyObject *args;
1355 PyObject *result;
1356 PyObject *it = lz->it;
1357
1358 assert(PyIter_Check(it));
1359 args = (*it->ob_type->tp_iternext)(it);
1360 if (args == NULL)
1361 return NULL;
1362 if (!PyTuple_CheckExact(args)) {
1363 Py_DECREF(args);
1364 PyErr_SetString(PyExc_TypeError,
1365 "iterator must return a tuple");
1366 return NULL;
1367 }
1368 result = PyObject_Call(lz->func, args, NULL);
1369 Py_DECREF(args);
1370 return result;
1371 }
1372
1373 PyDoc_STRVAR(starmap_doc,
1374 "starmap(function, sequence) --> starmap object\n\
1375 \n\
1376 Return an iterator whose values are returned from the function evaluated\n\
1377 with a argument tuple taken from the given sequence.");
1378
1379 static PyTypeObject starmap_type = {
1380 PyObject_HEAD_INIT(NULL)
1381 0, /* ob_size */
1382 "itertools.starmap", /* tp_name */
1383 sizeof(starmapobject), /* tp_basicsize */
1384 0, /* tp_itemsize */
1385 /* methods */
1386 (destructor)starmap_dealloc, /* tp_dealloc */
1387 0, /* tp_print */
1388 0, /* tp_getattr */
1389 0, /* tp_setattr */
1390 0, /* tp_compare */
1391 0, /* tp_repr */
1392 0, /* tp_as_number */
1393 0, /* tp_as_sequence */
1394 0, /* tp_as_mapping */
1395 0, /* tp_hash */
1396 0, /* tp_call */
1397 0, /* tp_str */
1398 PyObject_GenericGetAttr, /* tp_getattro */
1399 0, /* tp_setattro */
1400 0, /* tp_as_buffer */
1401 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1402 Py_TPFLAGS_BASETYPE, /* tp_flags */
1403 starmap_doc, /* tp_doc */
1404 (traverseproc)starmap_traverse, /* tp_traverse */
1405 0, /* tp_clear */
1406 0, /* tp_richcompare */
1407 0, /* tp_weaklistoffset */
1408 PyObject_SelfIter, /* tp_iter */
1409 (iternextfunc)starmap_next, /* tp_iternext */
1410 0, /* tp_methods */
1411 0, /* tp_members */
1412 0, /* tp_getset */
1413 0, /* tp_base */
1414 0, /* tp_dict */
1415 0, /* tp_descr_get */
1416 0, /* tp_descr_set */
1417 0, /* tp_dictoffset */
1418 0, /* tp_init */
1419 0, /* tp_alloc */
1420 starmap_new, /* tp_new */
1421 PyObject_GC_Del, /* tp_free */
1422 };
1423
1424
1425 /* imap object ************************************************************/
1426
1427 typedef struct {
1428 PyObject_HEAD
1429 PyObject *iters;
1430 PyObject *func;
1431 } imapobject;
1432
1433 static PyTypeObject imap_type;
1434
1435 static PyObject *
1436 imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1437 {
1438 PyObject *it, *iters, *func;
1439 imapobject *lz;
1440 Py_ssize_t numargs, i;
1441
1442 if (!_PyArg_NoKeywords("imap()", kwds))
1443 return NULL;
1444
1445 numargs = PyTuple_Size(args);
1446 if (numargs < 2) {
1447 PyErr_SetString(PyExc_TypeError,
1448 "imap() must have at least two arguments.");
1449 return NULL;
1450 }
1451
1452 iters = PyTuple_New(numargs-1);
1453 if (iters == NULL)
1454 return NULL;
1455
1456 for (i=1 ; i<numargs ; i++) {
1457 /* Get iterator. */
1458 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1459 if (it == NULL) {
1460 Py_DECREF(iters);
1461 return NULL;
1462 }
1463 PyTuple_SET_ITEM(iters, i-1, it);
1464 }
1465
1466 /* create imapobject structure */
1467 lz = (imapobject *)type->tp_alloc(type, 0);
1468 if (lz == NULL) {
1469 Py_DECREF(iters);
1470 return NULL;
1471 }
1472 lz->iters = iters;
1473 func = PyTuple_GET_ITEM(args, 0);
1474 Py_INCREF(func);
1475 lz->func = func;
1476
1477 return (PyObject *)lz;
1478 }
1479
1480 static void
1481 imap_dealloc(imapobject *lz)
1482 {
1483 PyObject_GC_UnTrack(lz);
1484 Py_XDECREF(lz->iters);
1485 Py_XDECREF(lz->func);
1486 lz->ob_type->tp_free(lz);
1487 }
1488
1489 static int
1490 imap_traverse(imapobject *lz, visitproc visit, void *arg)
1491 {
1492 Py_VISIT(lz->iters);
1493 Py_VISIT(lz->func);
1494 return 0;
1495 }
1496
1497 /*
1498 imap() is an iterator version of __builtins__.map() except that it does
1499 not have the None fill-in feature. That was intentionally left out for
1500 the following reasons:
1501
1502 1) Itertools are designed to be easily combined and chained together.
1503 Having all tools stop with the shortest input is a unifying principle
1504 that makes it easier to combine finite iterators (supplying data) with
1505 infinite iterators like count() and repeat() (for supplying sequential
1506 or constant arguments to a function).
1507
1508 2) In typical use cases for combining itertools, having one finite data
1509 supplier run out before another is likely to be an error condition which
1510 should not pass silently by automatically supplying None.
1511
1512 3) The use cases for automatic None fill-in are rare -- not many functions
1513 do something useful when a parameter suddenly switches type and becomes
1514 None.
1515
1516 4) If a need does arise, it can be met by __builtins__.map() or by
1517 writing: chain(iterable, repeat(None)).
1518
1519 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1520 */
1521
1522 static PyObject *
1523 imap_next(imapobject *lz)
1524 {
1525 PyObject *val;
1526 PyObject *argtuple;
1527 PyObject *result;
1528 Py_ssize_t numargs, i;
1529
1530 numargs = PyTuple_Size(lz->iters);
1531 argtuple = PyTuple_New(numargs);
1532 if (argtuple == NULL)
1533 return NULL;
1534
1535 for (i=0 ; i<numargs ; i++) {
1536 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1537 if (val == NULL) {
1538 Py_DECREF(argtuple);
1539 return NULL;
1540 }
1541 PyTuple_SET_ITEM(argtuple, i, val);
1542 }
1543 if (lz->func == Py_None)
1544 return argtuple;
1545 result = PyObject_Call(lz->func, argtuple, NULL);
1546 Py_DECREF(argtuple);
1547 return result;
1548 }
1549
1550 PyDoc_STRVAR(imap_doc,
1551 "imap(func, *iterables) --> imap object\n\
1552 \n\
1553 Make an iterator that computes the function using arguments from\n\
1554 each of the iterables. Like map() except that it returns\n\
1555 an iterator instead of a list and that it stops when the shortest\n\
1556 iterable is exhausted instead of filling in None for shorter\n\
1557 iterables.");
1558
1559 static PyTypeObject imap_type = {
1560 PyObject_HEAD_INIT(NULL)
1561 0, /* ob_size */
1562 "itertools.imap", /* tp_name */
1563 sizeof(imapobject), /* tp_basicsize */
1564 0, /* tp_itemsize */
1565 /* methods */
1566 (destructor)imap_dealloc, /* tp_dealloc */
1567 0, /* tp_print */
1568 0, /* tp_getattr */
1569 0, /* tp_setattr */
1570 0, /* tp_compare */
1571 0, /* tp_repr */
1572 0, /* tp_as_number */
1573 0, /* tp_as_sequence */
1574 0, /* tp_as_mapping */
1575 0, /* tp_hash */
1576 0, /* tp_call */
1577 0, /* tp_str */
1578 PyObject_GenericGetAttr, /* tp_getattro */
1579 0, /* tp_setattro */
1580 0, /* tp_as_buffer */
1581 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1582 Py_TPFLAGS_BASETYPE, /* tp_flags */
1583 imap_doc, /* tp_doc */
1584 (traverseproc)imap_traverse, /* tp_traverse */
1585 0, /* tp_clear */
1586 0, /* tp_richcompare */
1587 0, /* tp_weaklistoffset */
1588 PyObject_SelfIter, /* tp_iter */
1589 (iternextfunc)imap_next, /* tp_iternext */
1590 0, /* tp_methods */
1591 0, /* tp_members */
1592 0, /* tp_getset */
1593 0, /* tp_base */
1594 0, /* tp_dict */
1595 0, /* tp_descr_get */
1596 0, /* tp_descr_set */
1597 0, /* tp_dictoffset */
1598 0, /* tp_init */
1599 0, /* tp_alloc */
1600 imap_new, /* tp_new */
1601 PyObject_GC_Del, /* tp_free */
1602 };
1603
1604
1605 /* chain object ************************************************************/
1606
1607 typedef struct {
1608 PyObject_HEAD
1609 Py_ssize_t tuplesize;
1610 Py_ssize_t iternum; /* which iterator is active */
1611 PyObject *ittuple; /* tuple of iterators */
1612 } chainobject;
1613
1614 static PyTypeObject chain_type;
1615
1616 static PyObject *
1617 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1618 {
1619 chainobject *lz;
1620 Py_ssize_t tuplesize = PySequence_Length(args);
1621 Py_ssize_t i;
1622 PyObject *ittuple;
1623
1624 if (!_PyArg_NoKeywords("chain()", kwds))
1625 return NULL;
1626
1627 /* obtain iterators */
1628 assert(PyTuple_Check(args));
1629 ittuple = PyTuple_New(tuplesize);
1630 if (ittuple == NULL)
1631 return NULL;
1632 for (i=0; i < tuplesize; ++i) {
1633 PyObject *item = PyTuple_GET_ITEM(args, i);
1634 PyObject *it = PyObject_GetIter(item);
1635 if (it == NULL) {
1636 if (PyErr_ExceptionMatches(PyExc_TypeError))
1637 PyErr_Format(PyExc_TypeError,
1638 "chain argument #%zd must support iteration",
1639 i+1);
1640 Py_DECREF(ittuple);
1641 return NULL;
1642 }
1643 PyTuple_SET_ITEM(ittuple, i, it);
1644 }
1645
1646 /* create chainobject structure */
1647 lz = (chainobject *)type->tp_alloc(type, 0);
1648 if (lz == NULL) {
1649 Py_DECREF(ittuple);
1650 return NULL;
1651 }
1652
1653 lz->ittuple = ittuple;
1654 lz->iternum = 0;
1655 lz->tuplesize = tuplesize;
1656
1657 return (PyObject *)lz;
1658 }
1659
1660 static void
1661 chain_dealloc(chainobject *lz)
1662 {
1663 PyObject_GC_UnTrack(lz);
1664 Py_XDECREF(lz->ittuple);
1665 lz->ob_type->tp_free(lz);
1666 }
1667
1668 static int
1669 chain_traverse(chainobject *lz, visitproc visit, void *arg)
1670 {
1671 Py_VISIT(lz->ittuple);
1672 return 0;
1673 }
1674
1675 static PyObject *
1676 chain_next(chainobject *lz)
1677 {
1678 PyObject *it;
1679 PyObject *item;
1680
1681 while (lz->iternum < lz->tuplesize) {
1682 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1683 item = PyIter_Next(it);
1684 if (item != NULL)
1685 return item;
1686 if (PyErr_Occurred()) {
1687 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1688 PyErr_Clear();
1689 else
1690 return NULL;
1691 }
1692 lz->iternum++;
1693 }
1694 return NULL;
1695 }
1696
1697 PyDoc_STRVAR(chain_doc,
1698 "chain(*iterables) --> chain object\n\
1699 \n\
1700 Return a chain object whose .next() method returns elements from the\n\
1701 first iterable until it is exhausted, then elements from the next\n\
1702 iterable, until all of the iterables are exhausted.");
1703
1704 static PyTypeObject chain_type = {
1705 PyObject_HEAD_INIT(NULL)
1706 0, /* ob_size */
1707 "itertools.chain", /* tp_name */
1708 sizeof(chainobject), /* tp_basicsize */
1709 0, /* tp_itemsize */
1710 /* methods */
1711 (destructor)chain_dealloc, /* tp_dealloc */
1712 0, /* tp_print */
1713 0, /* tp_getattr */
1714 0, /* tp_setattr */
1715 0, /* tp_compare */
1716 0, /* tp_repr */
1717 0, /* tp_as_number */
1718 0, /* tp_as_sequence */
1719 0, /* tp_as_mapping */
1720 0, /* tp_hash */
1721 0, /* tp_call */
1722 0, /* tp_str */
1723 PyObject_GenericGetAttr, /* tp_getattro */
1724 0, /* tp_setattro */
1725 0, /* tp_as_buffer */
1726 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1727 Py_TPFLAGS_BASETYPE, /* tp_flags */
1728 chain_doc, /* tp_doc */
1729 (traverseproc)chain_traverse, /* tp_traverse */
1730 0, /* tp_clear */
1731 0, /* tp_richcompare */
1732 0, /* tp_weaklistoffset */
1733 PyObject_SelfIter, /* tp_iter */
1734 (iternextfunc)chain_next, /* tp_iternext */
1735 0, /* tp_methods */
1736 0, /* tp_members */
1737 0, /* tp_getset */
1738 0, /* tp_base */
1739 0, /* tp_dict */
1740 0, /* tp_descr_get */
1741 0, /* tp_descr_set */
1742 0, /* tp_dictoffset */
1743 0, /* tp_init */
1744 0, /* tp_alloc */
1745 chain_new, /* tp_new */
1746 PyObject_GC_Del, /* tp_free */
1747 };
1748
1749
1750 /* ifilter object ************************************************************/
1751
1752 typedef struct {
1753 PyObject_HEAD
1754 PyObject *func;
1755 PyObject *it;
1756 } ifilterobject;
1757
1758 static PyTypeObject ifilter_type;
1759
1760 static PyObject *
1761 ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1762 {
1763 PyObject *func, *seq;
1764 PyObject *it;
1765 ifilterobject *lz;
1766
1767 if (!_PyArg_NoKeywords("ifilter()", kwds))
1768 return NULL;
1769
1770 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
1771 return NULL;
1772
1773 /* Get iterator. */
1774 it = PyObject_GetIter(seq);
1775 if (it == NULL)
1776 return NULL;
1777
1778 /* create ifilterobject structure */
1779 lz = (ifilterobject *)type->tp_alloc(type, 0);
1780 if (lz == NULL) {
1781 Py_DECREF(it);
1782 return NULL;
1783 }
1784 Py_INCREF(func);
1785 lz->func = func;
1786 lz->it = it;
1787
1788 return (PyObject *)lz;
1789 }
1790
1791 static void
1792 ifilter_dealloc(ifilterobject *lz)
1793 {
1794 PyObject_GC_UnTrack(lz);
1795 Py_XDECREF(lz->func);
1796 Py_XDECREF(lz->it);
1797 lz->ob_type->tp_free(lz);
1798 }
1799
1800 static int
1801 ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1802 {
1803 Py_VISIT(lz->it);
1804 Py_VISIT(lz->func);
1805 return 0;
1806 }
1807
1808 static PyObject *
1809 ifilter_next(ifilterobject *lz)
1810 {
1811 PyObject *item;
1812 PyObject *it = lz->it;
1813 long ok;
1814 PyObject *(*iternext)(PyObject *);
1815
1816 assert(PyIter_Check(it));
1817 iternext = *it->ob_type->tp_iternext;
1818 for (;;) {
1819 item = iternext(it);
1820 if (item == NULL)
1821 return NULL;
1822
1823 if (lz->func == Py_None) {
1824 ok = PyObject_IsTrue(item);
1825 } else {
1826 PyObject *good;
1827 good = PyObject_CallFunctionObjArgs(lz->func,
1828 item, NULL);
1829 if (good == NULL) {
1830 Py_DECREF(item);
1831 return NULL;
1832 }
1833 ok = PyObject_IsTrue(good);
1834 Py_DECREF(good);
1835 }
1836 if (ok)
1837 return item;
1838 Py_DECREF(item);
1839 }
1840 }
1841
1842 PyDoc_STRVAR(ifilter_doc,
1843 "ifilter(function or None, sequence) --> ifilter object\n\
1844 \n\
1845 Return those items of sequence for which function(item) is true.\n\
1846 If function is None, return the items that are true.");
1847
1848 static PyTypeObject ifilter_type = {
1849 PyObject_HEAD_INIT(NULL)
1850 0, /* ob_size */
1851 "itertools.ifilter", /* tp_name */
1852 sizeof(ifilterobject), /* tp_basicsize */
1853 0, /* tp_itemsize */
1854 /* methods */
1855 (destructor)ifilter_dealloc, /* tp_dealloc */
1856 0, /* tp_print */
1857 0, /* tp_getattr */
1858 0, /* tp_setattr */
1859 0, /* tp_compare */
1860 0, /* tp_repr */
1861 0, /* tp_as_number */
1862 0, /* tp_as_sequence */
1863 0, /* tp_as_mapping */
1864 0, /* tp_hash */
1865 0, /* tp_call */
1866 0, /* tp_str */
1867 PyObject_GenericGetAttr, /* tp_getattro */
1868 0, /* tp_setattro */
1869 0, /* tp_as_buffer */
1870 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1871 Py_TPFLAGS_BASETYPE, /* tp_flags */
1872 ifilter_doc, /* tp_doc */
1873 (traverseproc)ifilter_traverse, /* tp_traverse */
1874 0, /* tp_clear */
1875 0, /* tp_richcompare */
1876 0, /* tp_weaklistoffset */
1877 PyObject_SelfIter, /* tp_iter */
1878 (iternextfunc)ifilter_next, /* tp_iternext */
1879 0, /* tp_methods */
1880 0, /* tp_members */
1881 0, /* tp_getset */
1882 0, /* tp_base */
1883 0, /* tp_dict */
1884 0, /* tp_descr_get */
1885 0, /* tp_descr_set */
1886 0, /* tp_dictoffset */
1887 0, /* tp_init */
1888 0, /* tp_alloc */
1889 ifilter_new, /* tp_new */
1890 PyObject_GC_Del, /* tp_free */
1891 };
1892
1893
1894 /* ifilterfalse object ************************************************************/
1895
1896 typedef struct {
1897 PyObject_HEAD
1898 PyObject *func;
1899 PyObject *it;
1900 } ifilterfalseobject;
1901
1902 static PyTypeObject ifilterfalse_type;
1903
1904 static PyObject *
1905 ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1906 {
1907 PyObject *func, *seq;
1908 PyObject *it;
1909 ifilterfalseobject *lz;
1910
1911 if (!_PyArg_NoKeywords("ifilterfalse()", kwds))
1912 return NULL;
1913
1914 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1915 return NULL;
1916
1917 /* Get iterator. */
1918 it = PyObject_GetIter(seq);
1919 if (it == NULL)
1920 return NULL;
1921
1922 /* create ifilterfalseobject structure */
1923 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1924 if (lz == NULL) {
1925 Py_DECREF(it);
1926 return NULL;
1927 }
1928 Py_INCREF(func);
1929 lz->func = func;
1930 lz->it = it;
1931
1932 return (PyObject *)lz;
1933 }
1934
1935 static void
1936 ifilterfalse_dealloc(ifilterfalseobject *lz)
1937 {
1938 PyObject_GC_UnTrack(lz);
1939 Py_XDECREF(lz->func);
1940 Py_XDECREF(lz->it);
1941 lz->ob_type->tp_free(lz);
1942 }
1943
1944 static int
1945 ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1946 {
1947 Py_VISIT(lz->it);
1948 Py_VISIT(lz->func);
1949 return 0;
1950 }
1951
1952 static PyObject *
1953 ifilterfalse_next(ifilterfalseobject *lz)
1954 {
1955 PyObject *item;
1956 PyObject *it = lz->it;
1957 long ok;
1958 PyObject *(*iternext)(PyObject *);
1959
1960 assert(PyIter_Check(it));
1961 iternext = *it->ob_type->tp_iternext;
1962 for (;;) {
1963 item = iternext(it);
1964 if (item == NULL)
1965 return NULL;
1966
1967 if (lz->func == Py_None) {
1968 ok = PyObject_IsTrue(item);
1969 } else {
1970 PyObject *good;
1971 good = PyObject_CallFunctionObjArgs(lz->func,
1972 item, NULL);
1973 if (good == NULL) {
1974 Py_DECREF(item);
1975 return NULL;
1976 }
1977 ok = PyObject_IsTrue(good);
1978 Py_DECREF(good);
1979 }
1980 if (!ok)
1981 return item;
1982 Py_DECREF(item);
1983 }
1984 }
1985
1986 PyDoc_STRVAR(ifilterfalse_doc,
1987 "ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1988 \n\
1989 Return those items of sequence for which function(item) is false.\n\
1990 If function is None, return the items that are false.");
1991
1992 static PyTypeObject ifilterfalse_type = {
1993 PyObject_HEAD_INIT(NULL)
1994 0, /* ob_size */
1995 "itertools.ifilterfalse", /* tp_name */
1996 sizeof(ifilterfalseobject), /* tp_basicsize */
1997 0, /* tp_itemsize */
1998 /* methods */
1999 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2000 0, /* tp_print */
2001 0, /* tp_getattr */
2002 0, /* tp_setattr */
2003 0, /* tp_compare */
2004 0, /* tp_repr */
2005 0, /* tp_as_number */
2006 0, /* tp_as_sequence */
2007 0, /* tp_as_mapping */
2008 0, /* tp_hash */
2009 0, /* tp_call */
2010 0, /* tp_str */
2011 PyObject_GenericGetAttr, /* tp_getattro */
2012 0, /* tp_setattro */
2013 0, /* tp_as_buffer */
2014 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2015 Py_TPFLAGS_BASETYPE, /* tp_flags */
2016 ifilterfalse_doc, /* tp_doc */
2017 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2018 0, /* tp_clear */
2019 0, /* tp_richcompare */
2020 0, /* tp_weaklistoffset */
2021 PyObject_SelfIter, /* tp_iter */
2022 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2023 0, /* tp_methods */
2024 0, /* tp_members */
2025 0, /* tp_getset */
2026 0, /* tp_base */
2027 0, /* tp_dict */
2028 0, /* tp_descr_get */
2029 0, /* tp_descr_set */
2030 0, /* tp_dictoffset */
2031 0, /* tp_init */
2032 0, /* tp_alloc */
2033 ifilterfalse_new, /* tp_new */
2034 PyObject_GC_Del, /* tp_free */
2035 };
2036
2037
2038 /* count object ************************************************************/
2039
2040 typedef struct {
2041 PyObject_HEAD
2042 Py_ssize_t cnt;
2043 } countobject;
2044
2045 static PyTypeObject count_type;
2046
2047 static PyObject *
2048 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2049 {
2050 countobject *lz;
2051 Py_ssize_t cnt = 0;
2052
2053 if (!_PyArg_NoKeywords("count()", kwds))
2054 return NULL;
2055
2056 if (!PyArg_ParseTuple(args, "|n:count", &cnt))
2057 return NULL;
2058
2059 /* create countobject structure */
2060 lz = (countobject *)PyObject_New(countobject, &count_type);
2061 if (lz == NULL)
2062 return NULL;
2063 lz->cnt = cnt;
2064
2065 return (PyObject *)lz;
2066 }
2067
2068 static PyObject *
2069 count_next(countobject *lz)
2070 {
2071 return PyInt_FromSize_t(lz->cnt++);
2072 }
2073
2074 static PyObject *
2075 count_repr(countobject *lz)
2076 {
2077 return PyString_FromFormat("count(%zd)", lz->cnt);
2078 }
2079
2080 PyDoc_STRVAR(count_doc,
2081 "count([firstval]) --> count object\n\
2082 \n\
2083 Return a count object whose .next() method returns consecutive\n\
2084 integers starting from zero or, if specified, from firstval.");
2085
2086 static PyTypeObject count_type = {
2087 PyObject_HEAD_INIT(NULL)
2088 0, /* ob_size */
2089 "itertools.count", /* tp_name */
2090 sizeof(countobject), /* tp_basicsize */
2091 0, /* tp_itemsize */
2092 /* methods */
2093 (destructor)PyObject_Del, /* tp_dealloc */
2094 0, /* tp_print */
2095 0, /* tp_getattr */
2096 0, /* tp_setattr */
2097 0, /* tp_compare */
2098 (reprfunc)count_repr, /* tp_repr */
2099 0, /* tp_as_number */
2100 0, /* tp_as_sequence */
2101 0, /* tp_as_mapping */
2102 0, /* tp_hash */
2103 0, /* tp_call */
2104 0, /* tp_str */
2105 PyObject_GenericGetAttr, /* tp_getattro */
2106 0, /* tp_setattro */
2107 0, /* tp_as_buffer */
2108 Py_TPFLAGS_DEFAULT, /* tp_flags */
2109 count_doc, /* tp_doc */
2110 0, /* tp_traverse */
2111 0, /* tp_clear */
2112 0, /* tp_richcompare */
2113 0, /* tp_weaklistoffset */
2114 PyObject_SelfIter, /* tp_iter */
2115 (iternextfunc)count_next, /* tp_iternext */
2116 0, /* tp_methods */
2117 0, /* tp_members */
2118 0, /* tp_getset */
2119 0, /* tp_base */
2120 0, /* tp_dict */
2121 0, /* tp_descr_get */
2122 0, /* tp_descr_set */
2123 0, /* tp_dictoffset */
2124 0, /* tp_init */
2125 0, /* tp_alloc */
2126 count_new, /* tp_new */
2127 };
2128
2129
2130 /* izip object ************************************************************/
2131
2132 #include "Python.h"
2133
2134 typedef struct {
2135 PyObject_HEAD
2136 Py_ssize_t tuplesize;
2137 PyObject *ittuple; /* tuple of iterators */
2138 PyObject *result;
2139 } izipobject;
2140
2141 static PyTypeObject izip_type;
2142
2143 static PyObject *
2144 izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2145 {
2146 izipobject *lz;
2147 Py_ssize_t i;
2148 PyObject *ittuple; /* tuple of iterators */
2149 PyObject *result;
2150 Py_ssize_t tuplesize = PySequence_Length(args);
2151
2152 if (!_PyArg_NoKeywords("izip()", kwds))
2153 return NULL;
2154
2155 /* args must be a tuple */
2156 assert(PyTuple_Check(args));
2157
2158 /* obtain iterators */
2159 ittuple = PyTuple_New(tuplesize);
2160 if (ittuple == NULL)
2161 return NULL;
2162 for (i=0; i < tuplesize; ++i) {
2163 PyObject *item = PyTuple_GET_ITEM(args, i);
2164 PyObject *it = PyObject_GetIter(item);
2165 if (it == NULL) {
2166 if (PyErr_ExceptionMatches(PyExc_TypeError))
2167 PyErr_Format(PyExc_TypeError,
2168 "izip argument #%zd must support iteration",
2169 i+1);
2170 Py_DECREF(ittuple);
2171 return NULL;
2172 }
2173 PyTuple_SET_ITEM(ittuple, i, it);
2174 }
2175
2176 /* create a result holder */
2177 result = PyTuple_New(tuplesize);
2178 if (result == NULL) {
2179 Py_DECREF(ittuple);
2180 return NULL;
2181 }
2182 for (i=0 ; i < tuplesize ; i++) {
2183 Py_INCREF(Py_None);
2184 PyTuple_SET_ITEM(result, i, Py_None);
2185 }
2186
2187 /* create izipobject structure */
2188 lz = (izipobject *)type->tp_alloc(type, 0);
2189 if (lz == NULL) {
2190 Py_DECREF(ittuple);
2191 Py_DECREF(result);
2192 return NULL;
2193 }
2194 lz->ittuple = ittuple;
2195 lz->tuplesize = tuplesize;
2196 lz->result = result;
2197
2198 return (PyObject *)lz;
2199 }
2200
2201 static void
2202 izip_dealloc(izipobject *lz)
2203 {
2204 PyObject_GC_UnTrack(lz);
2205 Py_XDECREF(lz->ittuple);
2206 Py_XDECREF(lz->result);
2207 lz->ob_type->tp_free(lz);
2208 }
2209
2210 static int
2211 izip_traverse(izipobject *lz, visitproc visit, void *arg)
2212 {
2213 Py_VISIT(lz->ittuple);
2214 Py_VISIT(lz->result);
2215 return 0;
2216 }
2217
2218 static PyObject *
2219 izip_next(izipobject *lz)
2220 {
2221 Py_ssize_t i;
2222 Py_ssize_t tuplesize = lz->tuplesize;
2223 PyObject *result = lz->result;
2224 PyObject *it;
2225 PyObject *item;
2226 PyObject *olditem;
2227
2228 if (tuplesize == 0)
2229 return NULL;
2230 if (result->ob_refcnt == 1) {
2231 Py_INCREF(result);
2232 for (i=0 ; i < tuplesize ; i++) {
2233 it = PyTuple_GET_ITEM(lz->ittuple, i);
2234 assert(PyIter_Check(it));
2235 item = (*it->ob_type->tp_iternext)(it);
2236 if (item == NULL) {
2237 Py_DECREF(result);
2238 return NULL;
2239 }
2240 olditem = PyTuple_GET_ITEM(result, i);
2241 PyTuple_SET_ITEM(result, i, item);
2242 Py_DECREF(olditem);
2243 }
2244 } else {
2245 result = PyTuple_New(tuplesize);
2246 if (result == NULL)
2247 return NULL;
2248 for (i=0 ; i < tuplesize ; i++) {
2249 it = PyTuple_GET_ITEM(lz->ittuple, i);
2250 assert(PyIter_Check(it));
2251 item = (*it->ob_type->tp_iternext)(it);
2252 if (item == NULL) {
2253 Py_DECREF(result);
2254 return NULL;
2255 }
2256 PyTuple_SET_ITEM(result, i, item);
2257 }
2258 }
2259 return result;
2260 }
2261
2262 PyDoc_STRVAR(izip_doc,
2263 "izip(iter1 [,iter2 [...]]) --> izip object\n\
2264 \n\
2265 Return a izip object whose .next() method returns a tuple where\n\
2266 the i-th element comes from the i-th iterable argument. The .next()\n\
2267 method continues until the shortest iterable in the argument sequence\n\
2268 is exhausted and then it raises StopIteration. Works like the zip()\n\
2269 function but consumes less memory by returning an iterator instead of\n\
2270 a list.");
2271
2272 static PyTypeObject izip_type = {
2273 PyObject_HEAD_INIT(NULL)
2274 0, /* ob_size */
2275 "itertools.izip", /* tp_name */
2276 sizeof(izipobject), /* tp_basicsize */
2277 0, /* tp_itemsize */
2278 /* methods */
2279 (destructor)izip_dealloc, /* tp_dealloc */
2280 0, /* tp_print */
2281 0, /* tp_getattr */
2282 0, /* tp_setattr */
2283 0, /* tp_compare */
2284 0, /* tp_repr */
2285 0, /* tp_as_number */
2286 0, /* tp_as_sequence */
2287 0, /* tp_as_mapping */
2288 0, /* tp_hash */
2289 0, /* tp_call */
2290 0, /* tp_str */
2291 PyObject_GenericGetAttr, /* tp_getattro */
2292 0, /* tp_setattro */
2293 0, /* tp_as_buffer */
2294 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2295 Py_TPFLAGS_BASETYPE, /* tp_flags */
2296 izip_doc, /* tp_doc */
2297 (traverseproc)izip_traverse, /* tp_traverse */
2298 0, /* tp_clear */
2299 0, /* tp_richcompare */
2300 0, /* tp_weaklistoffset */
2301 PyObject_SelfIter, /* tp_iter */
2302 (iternextfunc)izip_next, /* tp_iternext */
2303 0, /* tp_methods */
2304 0, /* tp_members */
2305 0, /* tp_getset */
2306 0, /* tp_base */
2307 0, /* tp_dict */
2308 0, /* tp_descr_get */
2309 0, /* tp_descr_set */
2310 0, /* tp_dictoffset */
2311 0, /* tp_init */
2312 0, /* tp_alloc */
2313 izip_new, /* tp_new */
2314 PyObject_GC_Del, /* tp_free */
2315 };
2316
2317
2318 /* repeat object ************************************************************/
2319
2320 typedef struct {
2321 PyObject_HEAD
2322 PyObject *element;
2323 Py_ssize_t cnt;
2324 } repeatobject;
2325
2326 static PyTypeObject repeat_type;
2327
2328 static PyObject *
2329 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2330 {
2331 repeatobject *ro;
2332 PyObject *element;
2333 Py_ssize_t cnt = -1;
2334
2335 if (!_PyArg_NoKeywords("repeat()", kwds))
2336 return NULL;
2337
2338 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
2339 return NULL;
2340
2341 if (PyTuple_Size(args) == 2 && cnt < 0)
2342 cnt = 0;
2343
2344 ro = (repeatobject *)type->tp_alloc(type, 0);
2345 if (ro == NULL)
2346 return NULL;
2347 Py_INCREF(element);
2348 ro->element = element;
2349 ro->cnt = cnt;
2350 return (PyObject *)ro;
2351 }
2352
2353 static void
2354 repeat_dealloc(repeatobject *ro)
2355 {
2356 PyObject_GC_UnTrack(ro);
2357 Py_XDECREF(ro->element);
2358 ro->ob_type->tp_free(ro);
2359 }
2360
2361 static int
2362 repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2363 {
2364 Py_VISIT(ro->element);
2365 return 0;
2366 }
2367
2368 static PyObject *
2369 repeat_next(repeatobject *ro)
2370 {
2371 if (ro->cnt == 0)
2372 return NULL;
2373 if (ro->cnt > 0)
2374 ro->cnt--;
2375 Py_INCREF(ro->element);
2376 return ro->element;
2377 }
2378
2379 static PyObject *
2380 repeat_repr(repeatobject *ro)
2381 {
2382 PyObject *result, *objrepr;
2383
2384 objrepr = PyObject_Repr(ro->element);
2385 if (objrepr == NULL)
2386 return NULL;
2387
2388 if (ro->cnt == -1)
2389 result = PyString_FromFormat("repeat(%s)",
2390 PyString_AS_STRING(objrepr));
2391 else
2392 result = PyString_FromFormat("repeat(%s, %zd)",
2393 PyString_AS_STRING(objrepr), ro->cnt);
2394 Py_DECREF(objrepr);
2395 return result;
2396 }
2397
2398 static PyObject *
2399 repeat_len(repeatobject *ro)
2400 {
2401 if (ro->cnt == -1) {
2402 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2403 return NULL;
2404 }
2405 return PyInt_FromSize_t(ro->cnt);
2406 }
2407
2408 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
2409
2410 static PyMethodDef repeat_methods[] = {
2411 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
2412 {NULL, NULL} /* sentinel */
2413 };
2414
2415 PyDoc_STRVAR(repeat_doc,
2416 "repeat(element [,times]) -> create an iterator which returns the element\n\
2417 for the specified number of times. If not specified, returns the element\n\
2418 endlessly.");
2419
2420 static PyTypeObject repeat_type = {
2421 PyObject_HEAD_INIT(NULL)
2422 0, /* ob_size */
2423 "itertools.repeat", /* tp_name */
2424 sizeof(repeatobject), /* tp_basicsize */
2425 0, /* tp_itemsize */
2426 /* methods */
2427 (destructor)repeat_dealloc, /* tp_dealloc */
2428 0, /* tp_print */
2429 0, /* tp_getattr */
2430 0, /* tp_setattr */
2431 0, /* tp_compare */
2432 (reprfunc)repeat_repr, /* tp_repr */
2433 0, /* tp_as_number */
2434 0, /* tp_as_sequence */
2435 0, /* tp_as_mapping */
2436 0, /* tp_hash */
2437 0, /* tp_call */
2438 0, /* tp_str */
2439 PyObject_GenericGetAttr, /* tp_getattro */
2440 0, /* tp_setattro */
2441 0, /* tp_as_buffer */
2442 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2443 Py_TPFLAGS_BASETYPE, /* tp_flags */
2444 repeat_doc, /* tp_doc */
2445 (traverseproc)repeat_traverse, /* tp_traverse */
2446 0, /* tp_clear */
2447 0, /* tp_richcompare */
2448 0, /* tp_weaklistoffset */
2449 PyObject_SelfIter, /* tp_iter */
2450 (iternextfunc)repeat_next, /* tp_iternext */
2451 repeat_methods, /* tp_methods */
2452 0, /* tp_members */
2453 0, /* tp_getset */
2454 0, /* tp_base */
2455 0, /* tp_dict */
2456 0, /* tp_descr_get */
2457 0, /* tp_descr_set */
2458 0, /* tp_dictoffset */
2459 0, /* tp_init */
2460 0, /* tp_alloc */
2461 repeat_new, /* tp_new */
2462 PyObject_GC_Del, /* tp_free */
2463 };
2464
2465
2466 /* module level code ********************************************************/
2467
2468 PyDoc_STRVAR(module_doc,
2469 "Functional tools for creating and using iterators.\n\
2470 \n\
2471 Infinite iterators:\n\
2472 count([n]) --> n, n+1, n+2, ...\n\
2473 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
2474 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
2475 \n\
2476 Iterators terminating on the shortest input sequence:\n\
2477 izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
2478 ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2479 ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
2480 islice(seq, [start,] stop [, step]) --> elements from\n\
2481 seq[start:stop:step]\n\
2482 imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2483 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
2484 tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
2485 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
2486 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2487 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
2488 groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
2489 ");
2490
2491
2492 static PyMethodDef module_methods[] = {
2493 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2494 {NULL, NULL} /* sentinel */
2495 };
2496
2497 PyMODINIT_FUNC
2498 inititertools(void)
2499 {
2500 int i;
2501 PyObject *m;
2502 char *name;
2503 PyTypeObject *typelist[] = {
2504 &cycle_type,
2505 &dropwhile_type,
2506 &takewhile_type,
2507 &islice_type,
2508 &starmap_type,
2509 &imap_type,
2510 &chain_type,
2511 &ifilter_type,
2512 &ifilterfalse_type,
2513 &count_type,
2514 &izip_type,
2515 &repeat_type,
2516 &groupby_type,
2517 NULL
2518 };
2519
2520 teedataobject_type.ob_type = &PyType_Type;
2521 m = Py_InitModule3("itertools", module_methods, module_doc);
2522 if (m == NULL)
2523 return;
2524
2525 for (i=0 ; typelist[i] != NULL ; i++) {
2526 if (PyType_Ready(typelist[i]) < 0)
2527 return;
2528 name = strchr(typelist[i]->tp_name, '.');
2529 assert (name != NULL);
2530 Py_INCREF(typelist[i]);
2531 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
2532 }
2533
2534 if (PyType_Ready(&teedataobject_type) < 0)
2535 return;
2536 if (PyType_Ready(&tee_type) < 0)
2537 return;
2538 if (PyType_Ready(&_grouper_type) < 0)
2539 return;
2540 }