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