?  search:
summary | log | graphiclog1 | graphiclog2 | commit | commitdiff | tree | refs | edit | fork
blame | history | raw | HEAD
Simplify code for itertools.product().
[python.git] / Modules / itertoolsmodule.c
blobbd5543cae77ea89b745db2ec9f9985f54aea38f4
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 PyObject *newargs = PySequence_Tuple(args);
1360 Py_DECREF(args);
1361 if (newargs == NULL)
1362 return NULL;
1363 args = newargs;
1364 }
1365 result = PyObject_Call(lz->func, args, NULL);
1366 Py_DECREF(args);
1367 return result;
1368 }
1369
1370 PyDoc_STRVAR(starmap_doc,
1371 "starmap(function, sequence) --> starmap object\n\
1372 \n\
1373 Return an iterator whose values are returned from the function evaluated\n\
1374 with a argument tuple taken from the given sequence.");
1375
1376 static PyTypeObject starmap_type = {
1377 PyVarObject_HEAD_INIT(NULL, 0)
1378 "itertools.starmap", /* tp_name */
1379 sizeof(starmapobject), /* tp_basicsize */
1380 0, /* tp_itemsize */
1381 /* methods */
1382 (destructor)starmap_dealloc, /* tp_dealloc */
1383 0, /* tp_print */
1384 0, /* tp_getattr */
1385 0, /* tp_setattr */
1386 0, /* tp_compare */
1387 0, /* tp_repr */
1388 0, /* tp_as_number */
1389 0, /* tp_as_sequence */
1390 0, /* tp_as_mapping */
1391 0, /* tp_hash */
1392 0, /* tp_call */
1393 0, /* tp_str */
1394 PyObject_GenericGetAttr, /* tp_getattro */
1395 0, /* tp_setattro */
1396 0, /* tp_as_buffer */
1397 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1398 Py_TPFLAGS_BASETYPE, /* tp_flags */
1399 starmap_doc, /* tp_doc */
1400 (traverseproc)starmap_traverse, /* tp_traverse */
1401 0, /* tp_clear */
1402 0, /* tp_richcompare */
1403 0, /* tp_weaklistoffset */
1404 PyObject_SelfIter, /* tp_iter */
1405 (iternextfunc)starmap_next, /* tp_iternext */
1406 0, /* tp_methods */
1407 0, /* tp_members */
1408 0, /* tp_getset */
1409 0, /* tp_base */
1410 0, /* tp_dict */
1411 0, /* tp_descr_get */
1412 0, /* tp_descr_set */
1413 0, /* tp_dictoffset */
1414 0, /* tp_init */
1415 0, /* tp_alloc */
1416 starmap_new, /* tp_new */
1417 PyObject_GC_Del, /* tp_free */
1418 };
1419
1420
1421 /* imap object ************************************************************/
1422
1423 typedef struct {
1424 PyObject_HEAD
1425 PyObject *iters;
1426 PyObject *func;
1427 } imapobject;
1428
1429 static PyTypeObject imap_type;
1430
1431 static PyObject *
1432 imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1433 {
1434 PyObject *it, *iters, *func;
1435 imapobject *lz;
1436 Py_ssize_t numargs, i;
1437
1438 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1439 return NULL;
1440
1441 numargs = PyTuple_Size(args);
1442 if (numargs < 2) {
1443 PyErr_SetString(PyExc_TypeError,
1444 "imap() must have at least two arguments.");
1445 return NULL;
1446 }
1447
1448 iters = PyTuple_New(numargs-1);
1449 if (iters == NULL)
1450 return NULL;
1451
1452 for (i=1 ; i<numargs ; i++) {
1453 /* Get iterator. */
1454 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1455 if (it == NULL) {
1456 Py_DECREF(iters);
1457 return NULL;
1458 }
1459 PyTuple_SET_ITEM(iters, i-1, it);
1460 }
1461
1462 /* create imapobject structure */
1463 lz = (imapobject *)type->tp_alloc(type, 0);
1464 if (lz == NULL) {
1465 Py_DECREF(iters);
1466 return NULL;
1467 }
1468 lz->iters = iters;
1469 func = PyTuple_GET_ITEM(args, 0);
1470 Py_INCREF(func);
1471 lz->func = func;
1472
1473 return (PyObject *)lz;
1474 }
1475
1476 static void
1477 imap_dealloc(imapobject *lz)
1478 {
1479 PyObject_GC_UnTrack(lz);
1480 Py_XDECREF(lz->iters);
1481 Py_XDECREF(lz->func);
1482 Py_TYPE(lz)->tp_free(lz);
1483 }
1484
1485 static int
1486 imap_traverse(imapobject *lz, visitproc visit, void *arg)
1487 {
1488 Py_VISIT(lz->iters);
1489 Py_VISIT(lz->func);
1490 return 0;
1491 }
1492
1493 /*
1494 imap() is an iterator version of __builtins__.map() except that it does
1495 not have the None fill-in feature. That was intentionally left out for
1496 the following reasons:
1497
1498 1) Itertools are designed to be easily combined and chained together.
1499 Having all tools stop with the shortest input is a unifying principle
1500 that makes it easier to combine finite iterators (supplying data) with
1501 infinite iterators like count() and repeat() (for supplying sequential
1502 or constant arguments to a function).
1503
1504 2) In typical use cases for combining itertools, having one finite data
1505 supplier run out before another is likely to be an error condition which
1506 should not pass silently by automatically supplying None.
1507
1508 3) The use cases for automatic None fill-in are rare -- not many functions
1509 do something useful when a parameter suddenly switches type and becomes
1510 None.
1511
1512 4) If a need does arise, it can be met by __builtins__.map() or by
1513 writing: chain(iterable, repeat(None)).
1514
1515 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1516 */
1517
1518 static PyObject *
1519 imap_next(imapobject *lz)
1520 {
1521 PyObject *val;
1522 PyObject *argtuple;
1523 PyObject *result;
1524 Py_ssize_t numargs, i;
1525
1526 numargs = PyTuple_Size(lz->iters);
1527 argtuple = PyTuple_New(numargs);
1528 if (argtuple == NULL)
1529 return NULL;
1530
1531 for (i=0 ; i<numargs ; i++) {
1532 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1533 if (val == NULL) {
1534 Py_DECREF(argtuple);
1535 return NULL;
1536 }
1537 PyTuple_SET_ITEM(argtuple, i, val);
1538 }
1539 if (lz->func == Py_None)
1540 return argtuple;
1541 result = PyObject_Call(lz->func, argtuple, NULL);
1542 Py_DECREF(argtuple);
1543 return result;
1544 }
1545
1546 PyDoc_STRVAR(imap_doc,
1547 "imap(func, *iterables) --> imap object\n\
1548 \n\
1549 Make an iterator that computes the function using arguments from\n\
1550 each of the iterables. Like map() except that it returns\n\
1551 an iterator instead of a list and that it stops when the shortest\n\
1552 iterable is exhausted instead of filling in None for shorter\n\
1553 iterables.");
1554
1555 static PyTypeObject imap_type = {
1556 PyVarObject_HEAD_INIT(NULL, 0)
1557 "itertools.imap", /* tp_name */
1558 sizeof(imapobject), /* tp_basicsize */
1559 0, /* tp_itemsize */
1560 /* methods */
1561 (destructor)imap_dealloc, /* tp_dealloc */
1562 0, /* tp_print */
1563 0, /* tp_getattr */
1564 0, /* tp_setattr */
1565 0, /* tp_compare */
1566 0, /* tp_repr */
1567 0, /* tp_as_number */
1568 0, /* tp_as_sequence */
1569 0, /* tp_as_mapping */
1570 0, /* tp_hash */
1571 0, /* tp_call */
1572 0, /* tp_str */
1573 PyObject_GenericGetAttr, /* tp_getattro */
1574 0, /* tp_setattro */
1575 0, /* tp_as_buffer */
1576 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1577 Py_TPFLAGS_BASETYPE, /* tp_flags */
1578 imap_doc, /* tp_doc */
1579 (traverseproc)imap_traverse, /* tp_traverse */
1580 0, /* tp_clear */
1581 0, /* tp_richcompare */
1582 0, /* tp_weaklistoffset */
1583 PyObject_SelfIter, /* tp_iter */
1584 (iternextfunc)imap_next, /* tp_iternext */
1585 0, /* tp_methods */
1586 0, /* tp_members */
1587 0, /* tp_getset */
1588 0, /* tp_base */
1589 0, /* tp_dict */
1590 0, /* tp_descr_get */
1591 0, /* tp_descr_set */
1592 0, /* tp_dictoffset */
1593 0, /* tp_init */
1594 0, /* tp_alloc */
1595 imap_new, /* tp_new */
1596 PyObject_GC_Del, /* tp_free */
1597 };
1598
1599
1600 /* chain object ************************************************************/
1601
1602 typedef struct {
1603 PyObject_HEAD
1604 PyObject *source; /* Iterator over input iterables */
1605 PyObject *active; /* Currently running input iterator */
1606 } chainobject;
1607
1608 static PyTypeObject chain_type;
1609
1610 static PyObject *
1611 chain_new_internal(PyTypeObject *type, PyObject *source)
1612 {
1613 chainobject *lz;
1614
1615 lz = (chainobject *)type->tp_alloc(type, 0);
1616 if (lz == NULL) {
1617 Py_DECREF(source);
1618 return NULL;
1619 }
1620
1621 lz->source = source;
1622 lz->active = NULL;
1623 return (PyObject *)lz;
1624 }
1625
1626 static PyObject *
1627 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1628 {
1629 PyObject *source;
1630
1631 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1632 return NULL;
1633
1634 source = PyObject_GetIter(args);
1635 if (source == NULL)
1636 return NULL;
1637
1638 return chain_new_internal(type, source);
1639 }
1640
1641 static PyObject *
1642 chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1643 {
1644 PyObject *source;
1645
1646 source = PyObject_GetIter(arg);
1647 if (source == NULL)
1648 return NULL;
1649
1650 return chain_new_internal(type, source);
1651 }
1652
1653 static void
1654 chain_dealloc(chainobject *lz)
1655 {
1656 PyObject_GC_UnTrack(lz);
1657 Py_XDECREF(lz->active);
1658 Py_XDECREF(lz->source);
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->source);
1666 Py_VISIT(lz->active);
1667 return 0;
1668 }
1669
1670 static PyObject *
1671 chain_next(chainobject *lz)
1672 {
1673 PyObject *item;
1674
1675 if (lz->source == NULL)
1676 return NULL; /* already stopped */
1677
1678 if (lz->active == NULL) {
1679 PyObject *iterable = PyIter_Next(lz->source);
1680 if (iterable == NULL) {
1681 Py_CLEAR(lz->source);
1682 return NULL; /* no more input sources */
1683 }
1684 lz->active = PyObject_GetIter(iterable);
1685 if (lz->active == NULL) {
1686 Py_DECREF(iterable);
1687 Py_CLEAR(lz->source);
1688 return NULL; /* input not iterable */
1689 }
1690 }
1691 item = PyIter_Next(lz->active);
1692 if (item != NULL)
1693 return item;
1694 if (PyErr_Occurred()) {
1695 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1696 PyErr_Clear();
1697 else
1698 return NULL; /* input raised an exception */
1699 }
1700 Py_CLEAR(lz->active);
1701 return chain_next(lz); /* recurse and use next active */
1702 }
1703
1704 PyDoc_STRVAR(chain_doc,
1705 "chain(*iterables) --> chain object\n\
1706 \n\
1707 Return a chain object whose .next() method returns elements from the\n\
1708 first iterable until it is exhausted, then elements from the next\n\
1709 iterable, until all of the iterables are exhausted.");
1710
1711 PyDoc_STRVAR(chain_from_iterable_doc,
1712 "chain.from_iterable(iterable) --> chain object\n\
1713 \n\
1714 Alternate chain() contructor taking a single iterable argument\n\
1715 that evaluates lazily.");
1716
1717 static PyMethodDef chain_methods[] = {
1718 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1719 chain_from_iterable_doc},
1720 {NULL, NULL} /* sentinel */
1721 };
1722
1723 static PyTypeObject chain_type = {
1724 PyVarObject_HEAD_INIT(NULL, 0)
1725 "itertools.chain", /* tp_name */
1726 sizeof(chainobject), /* tp_basicsize */
1727 0, /* tp_itemsize */
1728 /* methods */
1729 (destructor)chain_dealloc, /* tp_dealloc */
1730 0, /* tp_print */
1731 0, /* tp_getattr */
1732 0, /* tp_setattr */
1733 0, /* tp_compare */
1734 0, /* tp_repr */
1735 0, /* tp_as_number */
1736 0, /* tp_as_sequence */
1737 0, /* tp_as_mapping */
1738 0, /* tp_hash */
1739 0, /* tp_call */
1740 0, /* tp_str */
1741 PyObject_GenericGetAttr, /* tp_getattro */
1742 0, /* tp_setattro */
1743 0, /* tp_as_buffer */
1744 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1745 Py_TPFLAGS_BASETYPE, /* tp_flags */
1746 chain_doc, /* tp_doc */
1747 (traverseproc)chain_traverse, /* tp_traverse */
1748 0, /* tp_clear */
1749 0, /* tp_richcompare */
1750 0, /* tp_weaklistoffset */
1751 PyObject_SelfIter, /* tp_iter */
1752 (iternextfunc)chain_next, /* tp_iternext */
1753 chain_methods, /* tp_methods */
1754 0, /* tp_members */
1755 0, /* tp_getset */
1756 0, /* tp_base */
1757 0, /* tp_dict */
1758 0, /* tp_descr_get */
1759 0, /* tp_descr_set */
1760 0, /* tp_dictoffset */
1761 0, /* tp_init */
1762 0, /* tp_alloc */
1763 chain_new, /* tp_new */
1764 PyObject_GC_Del, /* tp_free */
1765 };
1766
1767
1768 /* product object ************************************************************/
1769
1770 typedef struct {
1771 PyObject_HEAD
1772 PyObject *pools; /* tuple of pool tuples */
1773 Py_ssize_t *indices; /* one index per pool */
1774 PyObject *result; /* most recently returned result tuple */
1775 int stopped; /* set to 1 when the product iterator is exhausted */
1776 } productobject;
1777
1778 static PyTypeObject product_type;
1779
1780 static PyObject *
1781 product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1782 {
1783 productobject *lz;
1784 Py_ssize_t nargs, npools, repeat=1;
1785 PyObject *pools = NULL;
1786 Py_ssize_t *indices = NULL;
1787 Py_ssize_t i;
1788
1789 if (kwds != NULL) {
1790 char *kwlist[] = {"repeat", 0};
1791 PyObject *tmpargs = PyTuple_New(0);
1792 if (tmpargs == NULL)
1793 return NULL;
1794 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1795 Py_DECREF(tmpargs);
1796 return NULL;
1797 }
1798 Py_DECREF(tmpargs);
1799 if (repeat < 0) {
1800 PyErr_SetString(PyExc_ValueError,
1801 "repeat argument cannot be negative");
1802 return NULL;
1803 }
1804 }
1805
1806 assert(PyTuple_Check(args));
1807 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1808 npools = nargs * repeat;
1809
1810 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1811 if (indices == NULL) {
1812 PyErr_NoMemory();
1813 goto error;
1814 }
1815
1816 pools = PyTuple_New(npools);
1817 if (pools == NULL)
1818 goto error;
1819
1820 for (i=0; i < nargs ; ++i) {
1821 PyObject *item = PyTuple_GET_ITEM(args, i);
1822 PyObject *pool = PySequence_Tuple(item);
1823 if (pool == NULL)
1824 goto error;
1825 PyTuple_SET_ITEM(pools, i, pool);
1826 indices[i] = 0;
1827 }
1828 for ( ; i < npools; ++i) {
1829 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1830 Py_INCREF(pool);
1831 PyTuple_SET_ITEM(pools, i, pool);
1832 indices[i] = 0;
1833 }
1834
1835 /* create productobject structure */
1836 lz = (productobject *)type->tp_alloc(type, 0);
1837 if (lz == NULL)
1838 goto error;
1839
1840 lz->pools = pools;
1841 lz->indices = indices;
1842 lz->result = NULL;
1843 lz->stopped = 0;
1844
1845 return (PyObject *)lz;
1846
1847 error:
1848 if (indices != NULL)
1849 PyMem_Free(indices);
1850 Py_XDECREF(pools);
1851 return NULL;
1852 }
1853
1854 static void
1855 product_dealloc(productobject *lz)
1856 {
1857 PyObject_GC_UnTrack(lz);
1858 Py_XDECREF(lz->pools);
1859 Py_XDECREF(lz->result);
1860 PyMem_Free(lz->indices);
1861 Py_TYPE(lz)->tp_free(lz);
1862 }
1863
1864 static int
1865 product_traverse(productobject *lz, visitproc visit, void *arg)
1866 {
1867 Py_VISIT(lz->pools);
1868 Py_VISIT(lz->result);
1869 return 0;
1870 }
1871
1872 static PyObject *
1873 product_next(productobject *lz)
1874 {
1875 PyObject *pool;
1876 PyObject *elem;
1877 PyObject *oldelem;
1878 PyObject *pools = lz->pools;
1879 PyObject *result = lz->result;
1880 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1881 Py_ssize_t i;
1882
1883 if (lz->stopped)
1884 return NULL;
1885
1886 if (result == NULL) {
1887 /* On the first pass, return an initial tuple filled with the
1888 first element from each pool. If any pool is empty, then
1889 whole product is empty and we're already done */
1890 if (npools == 0)
1891 goto empty;
1892 result = PyTuple_New(npools);
1893 if (result == NULL)
1894 goto empty;
1895 lz->result = result;
1896 for (i=0; i < npools; i++) {
1897 pool = PyTuple_GET_ITEM(pools, i);
1898 if (PyTuple_GET_SIZE(pool) == 0)
1899 goto empty;
1900 elem = PyTuple_GET_ITEM(pool, 0);
1901 Py_INCREF(elem);
1902 PyTuple_SET_ITEM(result, i, elem);
1903 }
1904 } else {
1905 Py_ssize_t *indices = lz->indices;
1906
1907 /* Copy the previous result tuple or re-use it if available */
1908 if (Py_REFCNT(result) > 1) {
1909 PyObject *old_result = result;
1910 result = PyTuple_New(npools);
1911 if (result == NULL)
1912 goto empty;
1913 lz->result = result;
1914 for (i=0; i < npools; i++) {
1915 elem = PyTuple_GET_ITEM(old_result, i);
1916 Py_INCREF(elem);
1917 PyTuple_SET_ITEM(result, i, elem);
1918 }
1919 Py_DECREF(old_result);
1920 }
1921 /* Now, we've got the only copy so we can update it in-place */
1922 assert (Py_REFCNT(result) == 1);
1923
1924 /* Update the pool indices right-to-left. Only advance to the
1925 next pool when the previous one rolls-over */
1926 for (i=npools-1 ; i >= 0 ; i--) {
1927 pool = PyTuple_GET_ITEM(pools, i);
1928 indices[i]++;
1929 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1930 /* Roll-over and advance to next pool */
1931 indices[i] = 0;
1932 elem = PyTuple_GET_ITEM(pool, 0);
1933 Py_INCREF(elem);
1934 oldelem = PyTuple_GET_ITEM(result, i);
1935 PyTuple_SET_ITEM(result, i, elem);
1936 Py_DECREF(oldelem);
1937 } else {
1938 /* No rollover. Just increment and stop here. */
1939 elem = PyTuple_GET_ITEM(pool, indices[i]);
1940 Py_INCREF(elem);
1941 oldelem = PyTuple_GET_ITEM(result, i);
1942 PyTuple_SET_ITEM(result, i, elem);
1943 Py_DECREF(oldelem);
1944 break;
1945 }
1946 }
1947
1948 /* If i is negative, then the indices have all rolled-over
1949 and we're done. */
1950 if (i < 0)
1951 goto empty;
1952 }
1953
1954 Py_INCREF(result);
1955 return result;
1956
1957 empty:
1958 lz->stopped = 1;
1959 return NULL;
1960 }
1961
1962 PyDoc_STRVAR(product_doc,
1963 "product(*iterables) --> product object\n\
1964 \n\
1965 Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
1966 For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1967 The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1968 cycle in a manner similar to an odometer (with the rightmost element changing\n\
1969 on every iteration).\n\n\
1970 product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1971 product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1972
1973 static PyTypeObject product_type = {
1974 PyVarObject_HEAD_INIT(NULL, 0)
1975 "itertools.product", /* tp_name */
1976 sizeof(productobject), /* tp_basicsize */
1977 0, /* tp_itemsize */
1978 /* methods */
1979 (destructor)product_dealloc, /* tp_dealloc */
1980 0, /* tp_print */
1981 0, /* tp_getattr */
1982 0, /* tp_setattr */
1983 0, /* tp_compare */
1984 0, /* tp_repr */
1985 0, /* tp_as_number */
1986 0, /* tp_as_sequence */
1987 0, /* tp_as_mapping */
1988 0, /* tp_hash */
1989 0, /* tp_call */
1990 0, /* tp_str */
1991 PyObject_GenericGetAttr, /* tp_getattro */
1992 0, /* tp_setattro */
1993 0, /* tp_as_buffer */
1994 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1995 Py_TPFLAGS_BASETYPE, /* tp_flags */
1996 product_doc, /* tp_doc */
1997 (traverseproc)product_traverse, /* tp_traverse */
1998 0, /* tp_clear */
1999 0, /* tp_richcompare */
2000 0, /* tp_weaklistoffset */
2001 PyObject_SelfIter, /* tp_iter */
2002 (iternextfunc)product_next, /* tp_iternext */
2003 0, /* tp_methods */
2004 0, /* tp_members */
2005 0, /* tp_getset */
2006 0, /* tp_base */
2007 0, /* tp_dict */
2008 0, /* tp_descr_get */
2009 0, /* tp_descr_set */
2010 0, /* tp_dictoffset */
2011 0, /* tp_init */
2012 0, /* tp_alloc */
2013 product_new, /* tp_new */
2014 PyObject_GC_Del, /* tp_free */
2015 };
2016
2017
2018 /* combinations object ************************************************************/
2019
2020 typedef struct {
2021 PyObject_HEAD
2022 PyObject *pool; /* input converted to a tuple */
2023 Py_ssize_t *indices; /* one index per result element */
2024 PyObject *result; /* most recently returned result tuple */
2025 Py_ssize_t r; /* size of result tuple */
2026 int stopped; /* set to 1 when the combinations iterator is exhausted */
2027 } combinationsobject;
2028
2029 static PyTypeObject combinations_type;
2030
2031 static PyObject *
2032 combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2033 {
2034 combinationsobject *co;
2035 Py_ssize_t n;
2036 Py_ssize_t r;
2037 PyObject *pool = NULL;
2038 PyObject *iterable = NULL;
2039 Py_ssize_t *indices = NULL;
2040 Py_ssize_t i;
2041 static char *kwargs[] = {"iterable", "r", NULL};
2042
2043 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2044 &iterable, &r))
2045 return NULL;
2046
2047 pool = PySequence_Tuple(iterable);
2048 if (pool == NULL)
2049 goto error;
2050 n = PyTuple_GET_SIZE(pool);
2051 if (r < 0) {
2052 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2053 goto error;
2054 }
2055 if (r > n) {
2056 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
2057 goto error;
2058 }
2059
2060 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2061 if (indices == NULL) {
2062 PyErr_NoMemory();
2063 goto error;
2064 }
2065
2066 for (i=0 ; i<r ; i++)
2067 indices[i] = i;
2068
2069 /* create combinationsobject structure */
2070 co = (combinationsobject *)type->tp_alloc(type, 0);
2071 if (co == NULL)
2072 goto error;
2073
2074 co->pool = pool;
2075 co->indices = indices;
2076 co->result = NULL;
2077 co->r = r;
2078 co->stopped = 0;
2079
2080 return (PyObject *)co;
2081
2082 error:
2083 if (indices != NULL)
2084 PyMem_Free(indices);
2085 Py_XDECREF(pool);
2086 return NULL;
2087 }
2088
2089 static void
2090 combinations_dealloc(combinationsobject *co)
2091 {
2092 PyObject_GC_UnTrack(co);
2093 Py_XDECREF(co->pool);
2094 Py_XDECREF(co->result);
2095 PyMem_Free(co->indices);
2096 Py_TYPE(co)->tp_free(co);
2097 }
2098
2099 static int
2100 combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2101 {
2102 Py_VISIT(co->pool);
2103 Py_VISIT(co->result);
2104 return 0;
2105 }
2106
2107 static PyObject *
2108 combinations_next(combinationsobject *co)
2109 {
2110 PyObject *elem;
2111 PyObject *oldelem;
2112 PyObject *pool = co->pool;
2113 Py_ssize_t *indices = co->indices;
2114 PyObject *result = co->result;
2115 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2116 Py_ssize_t r = co->r;
2117 Py_ssize_t i, j, index;
2118
2119 if (co->stopped)
2120 return NULL;
2121
2122 if (result == NULL) {
2123 /* On the first pass, initialize result tuple using the indices */
2124 result = PyTuple_New(r);
2125 if (result == NULL)
2126 goto empty;
2127 co->result = result;
2128 for (i=0; i<r ; i++) {
2129 index = indices[i];
2130 elem = PyTuple_GET_ITEM(pool, index);
2131 Py_INCREF(elem);
2132 PyTuple_SET_ITEM(result, i, elem);
2133 }
2134 } else {
2135 /* Copy the previous result tuple or re-use it if available */
2136 if (Py_REFCNT(result) > 1) {
2137 PyObject *old_result = result;
2138 result = PyTuple_New(r);
2139 if (result == NULL)
2140 goto empty;
2141 co->result = result;
2142 for (i=0; i<r ; i++) {
2143 elem = PyTuple_GET_ITEM(old_result, i);
2144 Py_INCREF(elem);
2145 PyTuple_SET_ITEM(result, i, elem);
2146 }
2147 Py_DECREF(old_result);
2148 }
2149 /* Now, we've got the only copy so we can update it in-place
2150 * CPython's empty tuple is a singleton and cached in
2151 * PyTuple's freelist.
2152 */
2153 assert(r == 0 || Py_REFCNT(result) == 1);
2154
2155 /* Scan indices right-to-left until finding one that is not
2156 at its maximum (i + n - r). */
2157 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2158 ;
2159
2160 /* If i is negative, then the indices are all at
2161 their maximum value and we're done. */
2162 if (i < 0)
2163 goto empty;
2164
2165 /* Increment the current index which we know is not at its
2166 maximum. Then move back to the right setting each index
2167 to its lowest possible value (one higher than the index
2168 to its left -- this maintains the sort order invariant). */
2169 indices[i]++;
2170 for (j=i+1 ; j<r ; j++)
2171 indices[j] = indices[j-1] + 1;
2172
2173 /* Update the result tuple for the new indices
2174 starting with i, the leftmost index that changed */
2175 for ( ; i<r ; i++) {
2176 index = indices[i];
2177 elem = PyTuple_GET_ITEM(pool, index);
2178 Py_INCREF(elem);
2179 oldelem = PyTuple_GET_ITEM(result, i);
2180 PyTuple_SET_ITEM(result, i, elem);
2181 Py_DECREF(oldelem);
2182 }
2183 }
2184
2185 Py_INCREF(result);
2186 return result;
2187
2188 empty:
2189 co->stopped = 1;
2190 return NULL;
2191 }
2192
2193 PyDoc_STRVAR(combinations_doc,
2194 "combinations(iterables) --> combinations object\n\
2195 \n\
2196 Return successive r-length combinations of elements in the iterable.\n\n\
2197 combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2198
2199 static PyTypeObject combinations_type = {
2200 PyVarObject_HEAD_INIT(NULL, 0)
2201 "itertools.combinations", /* tp_name */
2202 sizeof(combinationsobject), /* tp_basicsize */
2203 0, /* tp_itemsize */
2204 /* methods */
2205 (destructor)combinations_dealloc, /* tp_dealloc */
2206 0, /* tp_print */
2207 0, /* tp_getattr */
2208 0, /* tp_setattr */
2209 0, /* tp_compare */
2210 0, /* tp_repr */
2211 0, /* tp_as_number */
2212 0, /* tp_as_sequence */
2213 0, /* tp_as_mapping */
2214 0, /* tp_hash */
2215 0, /* tp_call */
2216 0, /* tp_str */
2217 PyObject_GenericGetAttr, /* tp_getattro */
2218 0, /* tp_setattro */
2219 0, /* tp_as_buffer */
2220 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2221 Py_TPFLAGS_BASETYPE, /* tp_flags */
2222 combinations_doc, /* tp_doc */
2223 (traverseproc)combinations_traverse, /* tp_traverse */
2224 0, /* tp_clear */
2225 0, /* tp_richcompare */
2226 0, /* tp_weaklistoffset */
2227 PyObject_SelfIter, /* tp_iter */
2228 (iternextfunc)combinations_next, /* tp_iternext */
2229 0, /* tp_methods */
2230 0, /* tp_members */
2231 0, /* tp_getset */
2232 0, /* tp_base */
2233 0, /* tp_dict */
2234 0, /* tp_descr_get */
2235 0, /* tp_descr_set */
2236 0, /* tp_dictoffset */
2237 0, /* tp_init */
2238 0, /* tp_alloc */
2239 combinations_new, /* tp_new */
2240 PyObject_GC_Del, /* tp_free */
2241 };
2242
2243
2244 /* ifilter object ************************************************************/
2245
2246 typedef struct {
2247 PyObject_HEAD
2248 PyObject *func;
2249 PyObject *it;
2250 } ifilterobject;
2251
2252 static PyTypeObject ifilter_type;
2253
2254 static PyObject *
2255 ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2256 {
2257 PyObject *func, *seq;
2258 PyObject *it;
2259 ifilterobject *lz;
2260
2261 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2262 return NULL;
2263
2264 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2265 return NULL;
2266
2267 /* Get iterator. */
2268 it = PyObject_GetIter(seq);
2269 if (it == NULL)
2270 return NULL;
2271
2272 /* create ifilterobject structure */
2273 lz = (ifilterobject *)type->tp_alloc(type, 0);
2274 if (lz == NULL) {
2275 Py_DECREF(it);
2276 return NULL;
2277 }
2278 Py_INCREF(func);
2279 lz->func = func;
2280 lz->it = it;
2281
2282 return (PyObject *)lz;
2283 }
2284
2285 static void
2286 ifilter_dealloc(ifilterobject *lz)
2287 {
2288 PyObject_GC_UnTrack(lz);
2289 Py_XDECREF(lz->func);
2290 Py_XDECREF(lz->it);
2291 Py_TYPE(lz)->tp_free(lz);
2292 }
2293
2294 static int
2295 ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2296 {
2297 Py_VISIT(lz->it);
2298 Py_VISIT(lz->func);
2299 return 0;
2300 }
2301
2302 static PyObject *
2303 ifilter_next(ifilterobject *lz)
2304 {
2305 PyObject *item;
2306 PyObject *it = lz->it;
2307 long ok;
2308 PyObject *(*iternext)(PyObject *);
2309
2310 assert(PyIter_Check(it));
2311 iternext = *Py_TYPE(it)->tp_iternext;
2312 for (;;) {
2313 item = iternext(it);
2314 if (item == NULL)
2315 return NULL;
2316
2317 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2318 ok = PyObject_IsTrue(item);
2319 } else {
2320 PyObject *good;
2321 good = PyObject_CallFunctionObjArgs(lz->func,
2322 item, NULL);
2323 if (good == NULL) {
2324 Py_DECREF(item);
2325 return NULL;
2326 }
2327 ok = PyObject_IsTrue(good);
2328 Py_DECREF(good);
2329 }
2330 if (ok)
2331 return item;
2332 Py_DECREF(item);
2333 }
2334 }
2335
2336 PyDoc_STRVAR(ifilter_doc,
2337 "ifilter(function or None, sequence) --> ifilter object\n\
2338 \n\
2339 Return those items of sequence for which function(item) is true.\n\
2340 If function is None, return the items that are true.");
2341
2342 static PyTypeObject ifilter_type = {
2343 PyVarObject_HEAD_INIT(NULL, 0)
2344 "itertools.ifilter", /* tp_name */
2345 sizeof(ifilterobject), /* tp_basicsize */
2346 0, /* tp_itemsize */
2347 /* methods */
2348 (destructor)ifilter_dealloc, /* tp_dealloc */
2349 0, /* tp_print */
2350 0, /* tp_getattr */
2351 0, /* tp_setattr */
2352 0, /* tp_compare */
2353 0, /* tp_repr */
2354 0, /* tp_as_number */
2355 0, /* tp_as_sequence */
2356 0, /* tp_as_mapping */
2357 0, /* tp_hash */
2358 0, /* tp_call */
2359 0, /* tp_str */
2360 PyObject_GenericGetAttr, /* tp_getattro */
2361 0, /* tp_setattro */
2362 0, /* tp_as_buffer */
2363 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2364 Py_TPFLAGS_BASETYPE, /* tp_flags */
2365 ifilter_doc, /* tp_doc */
2366 (traverseproc)ifilter_traverse, /* tp_traverse */
2367 0, /* tp_clear */
2368 0, /* tp_richcompare */
2369 0, /* tp_weaklistoffset */
2370 PyObject_SelfIter, /* tp_iter */
2371 (iternextfunc)ifilter_next, /* tp_iternext */
2372 0, /* tp_methods */
2373 0, /* tp_members */
2374 0, /* tp_getset */
2375 0, /* tp_base */
2376 0, /* tp_dict */
2377 0, /* tp_descr_get */
2378 0, /* tp_descr_set */
2379 0, /* tp_dictoffset */
2380 0, /* tp_init */
2381 0, /* tp_alloc */
2382 ifilter_new, /* tp_new */
2383 PyObject_GC_Del, /* tp_free */
2384 };
2385
2386
2387 /* ifilterfalse object ************************************************************/
2388
2389 typedef struct {
2390 PyObject_HEAD
2391 PyObject *func;
2392 PyObject *it;
2393 } ifilterfalseobject;
2394
2395 static PyTypeObject ifilterfalse_type;
2396
2397 static PyObject *
2398 ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2399 {
2400 PyObject *func, *seq;
2401 PyObject *it;
2402 ifilterfalseobject *lz;
2403
2404 if (type == &ifilterfalse_type &&
2405 !_PyArg_NoKeywords("ifilterfalse()", kwds))
2406 return NULL;
2407
2408 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2409 return NULL;
2410
2411 /* Get iterator. */
2412 it = PyObject_GetIter(seq);
2413 if (it == NULL)
2414 return NULL;
2415
2416 /* create ifilterfalseobject structure */
2417 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
2418 if (lz == NULL) {
2419 Py_DECREF(it);
2420 return NULL;
2421 }
2422 Py_INCREF(func);
2423 lz->func = func;
2424 lz->it = it;
2425
2426 return (PyObject *)lz;
2427 }
2428
2429 static void
2430 ifilterfalse_dealloc(ifilterfalseobject *lz)
2431 {
2432 PyObject_GC_UnTrack(lz);
2433 Py_XDECREF(lz->func);
2434 Py_XDECREF(lz->it);
2435 Py_TYPE(lz)->tp_free(lz);
2436 }
2437
2438 static int
2439 ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
2440 {
2441 Py_VISIT(lz->it);
2442 Py_VISIT(lz->func);
2443 return 0;
2444 }
2445
2446 static PyObject *
2447 ifilterfalse_next(ifilterfalseobject *lz)
2448 {
2449 PyObject *item;
2450 PyObject *it = lz->it;
2451 long ok;
2452 PyObject *(*iternext)(PyObject *);
2453
2454 assert(PyIter_Check(it));
2455 iternext = *Py_TYPE(it)->tp_iternext;
2456 for (;;) {
2457 item = iternext(it);
2458 if (item == NULL)
2459 return NULL;
2460
2461 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2462 ok = PyObject_IsTrue(item);
2463 } else {
2464 PyObject *good;
2465 good = PyObject_CallFunctionObjArgs(lz->func,
2466 item, NULL);
2467 if (good == NULL) {
2468 Py_DECREF(item);
2469 return NULL;
2470 }
2471 ok = PyObject_IsTrue(good);
2472 Py_DECREF(good);
2473 }
2474 if (!ok)
2475 return item;
2476 Py_DECREF(item);
2477 }
2478 }
2479
2480 PyDoc_STRVAR(ifilterfalse_doc,
2481 "ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2482 \n\
2483 Return those items of sequence for which function(item) is false.\n\
2484 If function is None, return the items that are false.");
2485
2486 static PyTypeObject ifilterfalse_type = {
2487 PyVarObject_HEAD_INIT(NULL, 0)
2488 "itertools.ifilterfalse", /* tp_name */
2489 sizeof(ifilterfalseobject), /* tp_basicsize */
2490 0, /* tp_itemsize */
2491 /* methods */
2492 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2493 0, /* tp_print */
2494 0, /* tp_getattr */
2495 0, /* tp_setattr */
2496 0, /* tp_compare */
2497 0, /* tp_repr */
2498 0, /* tp_as_number */
2499 0, /* tp_as_sequence */
2500 0, /* tp_as_mapping */
2501 0, /* tp_hash */
2502 0, /* tp_call */
2503 0, /* tp_str */
2504 PyObject_GenericGetAttr, /* tp_getattro */
2505 0, /* tp_setattro */
2506 0, /* tp_as_buffer */
2507 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2508 Py_TPFLAGS_BASETYPE, /* tp_flags */
2509 ifilterfalse_doc, /* tp_doc */
2510 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2511 0, /* tp_clear */
2512 0, /* tp_richcompare */
2513 0, /* tp_weaklistoffset */
2514 PyObject_SelfIter, /* tp_iter */
2515 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2516 0, /* tp_methods */
2517 0, /* tp_members */
2518 0, /* tp_getset */
2519 0, /* tp_base */
2520 0, /* tp_dict */
2521 0, /* tp_descr_get */
2522 0, /* tp_descr_set */
2523 0, /* tp_dictoffset */
2524 0, /* tp_init */
2525 0, /* tp_alloc */
2526 ifilterfalse_new, /* tp_new */
2527 PyObject_GC_Del, /* tp_free */
2528 };
2529
2530
2531 /* count object ************************************************************/
2532
2533 typedef struct {
2534 PyObject_HEAD
2535 Py_ssize_t cnt;
2536 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
2537 } countobject;
2538
2539 static PyTypeObject count_type;
2540
2541 static PyObject *
2542 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2543 {
2544 countobject *lz;
2545 Py_ssize_t cnt = 0;
2546 PyObject *cnt_arg = NULL;
2547 PyObject *long_cnt = NULL;
2548
2549 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
2550 return NULL;
2551
2552 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
2553 return NULL;
2554
2555 if (cnt_arg != NULL) {
2556 cnt = PyInt_AsSsize_t(cnt_arg);
2557 if (cnt == -1 && PyErr_Occurred()) {
2558 PyErr_Clear();
2559 if (!PyLong_Check(cnt_arg)) {
2560 PyErr_SetString(PyExc_TypeError, "an integer is required");
2561 return NULL;
2562 }
2563 long_cnt = cnt_arg;
2564 Py_INCREF(long_cnt);
2565 cnt = PY_SSIZE_T_MAX;
2566 }
2567 }
2568
2569 /* create countobject structure */
2570 lz = (countobject *)PyObject_New(countobject, &count_type);
2571 if (lz == NULL) {
2572 Py_XDECREF(long_cnt);
2573 return NULL;
2574 }
2575 lz->cnt = cnt;
2576 lz->long_cnt = long_cnt;
2577
2578 return (PyObject *)lz;
2579 }
2580
2581 static void
2582 count_dealloc(countobject *lz)
2583 {
2584 Py_XDECREF(lz->long_cnt);
2585 PyObject_Del(lz);
2586 }
2587
2588 static PyObject *
2589 count_nextlong(countobject *lz)
2590 {
2591 static PyObject *one = NULL;
2592 PyObject *cnt;
2593 PyObject *stepped_up;
2594
2595 if (lz->long_cnt == NULL) {
2596 lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
2597 if (lz->long_cnt == NULL)
2598 return NULL;
2599 }
2600 if (one == NULL) {
2601 one = PyInt_FromLong(1);
2602 if (one == NULL)
2603 return NULL;
2604 }
2605 cnt = lz->long_cnt;
2606 assert(cnt != NULL);
2607 stepped_up = PyNumber_Add(cnt, one);
2608 if (stepped_up == NULL)
2609 return NULL;
2610 lz->long_cnt = stepped_up;
2611 return cnt;
2612 }
2613
2614 static PyObject *
2615 count_next(countobject *lz)
2616 {
2617 if (lz->cnt == PY_SSIZE_T_MAX)
2618 return count_nextlong(lz);
2619 return PyInt_FromSsize_t(lz->cnt++);
2620 }
2621
2622 static PyObject *
2623 count_repr(countobject *lz)
2624 {
2625 PyObject *cnt_repr;
2626 PyObject *result;
2627
2628 if (lz->cnt != PY_SSIZE_T_MAX)
2629 return PyString_FromFormat("count(%zd)", lz->cnt);
2630
2631 cnt_repr = PyObject_Repr(lz->long_cnt);
2632 if (cnt_repr == NULL)
2633 return NULL;
2634 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
2635 Py_DECREF(cnt_repr);
2636 return result;
2637 }
2638
2639 PyDoc_STRVAR(count_doc,
2640 "count([firstval]) --> count object\n\
2641 \n\
2642 Return a count object whose .next() method returns consecutive\n\
2643 integers starting from zero or, if specified, from firstval.");
2644
2645 static PyTypeObject count_type = {
2646 PyVarObject_HEAD_INIT(NULL, 0)
2647 "itertools.count", /* tp_name */
2648 sizeof(countobject), /* tp_basicsize */
2649 0, /* tp_itemsize */
2650 /* methods */
2651 (destructor)count_dealloc, /* tp_dealloc */
2652 0, /* tp_print */
2653 0, /* tp_getattr */
2654 0, /* tp_setattr */
2655 0, /* tp_compare */
2656 (reprfunc)count_repr, /* tp_repr */
2657 0, /* tp_as_number */
2658 0, /* tp_as_sequence */
2659 0, /* tp_as_mapping */
2660 0, /* tp_hash */
2661 0, /* tp_call */
2662 0, /* tp_str */
2663 PyObject_GenericGetAttr, /* tp_getattro */
2664 0, /* tp_setattro */
2665 0, /* tp_as_buffer */
2666 Py_TPFLAGS_DEFAULT, /* tp_flags */
2667 count_doc, /* tp_doc */
2668 0, /* tp_traverse */
2669 0, /* tp_clear */
2670 0, /* tp_richcompare */
2671 0, /* tp_weaklistoffset */
2672 PyObject_SelfIter, /* tp_iter */
2673 (iternextfunc)count_next, /* tp_iternext */
2674 0, /* tp_methods */
2675 0, /* tp_members */
2676 0, /* tp_getset */
2677 0, /* tp_base */
2678 0, /* tp_dict */
2679 0, /* tp_descr_get */
2680 0, /* tp_descr_set */
2681 0, /* tp_dictoffset */
2682 0, /* tp_init */
2683 0, /* tp_alloc */
2684 count_new, /* tp_new */
2685 };
2686
2687
2688 /* izip object ************************************************************/
2689
2690 #include "Python.h"
2691
2692 typedef struct {
2693 PyObject_HEAD
2694 Py_ssize_t tuplesize;
2695 PyObject *ittuple; /* tuple of iterators */
2696 PyObject *result;
2697 } izipobject;
2698
2699 static PyTypeObject izip_type;
2700
2701 static PyObject *
2702 izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2703 {
2704 izipobject *lz;
2705 Py_ssize_t i;
2706 PyObject *ittuple; /* tuple of iterators */
2707 PyObject *result;
2708 Py_ssize_t tuplesize = PySequence_Length(args);
2709
2710 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
2711 return NULL;
2712
2713 /* args must be a tuple */
2714 assert(PyTuple_Check(args));
2715
2716 /* obtain iterators */
2717 ittuple = PyTuple_New(tuplesize);
2718 if (ittuple == NULL)
2719 return NULL;
2720 for (i=0; i < tuplesize; ++i) {
2721 PyObject *item = PyTuple_GET_ITEM(args, i);
2722 PyObject *it = PyObject_GetIter(item);
2723 if (it == NULL) {
2724 if (PyErr_ExceptionMatches(PyExc_TypeError))
2725 PyErr_Format(PyExc_TypeError,
2726 "izip argument #%zd must support iteration",
2727 i+1);
2728 Py_DECREF(ittuple);
2729 return NULL;
2730 }
2731 PyTuple_SET_ITEM(ittuple, i, it);
2732 }
2733
2734 /* create a result holder */
2735 result = PyTuple_New(tuplesize);
2736 if (result == NULL) {
2737 Py_DECREF(ittuple);
2738 return NULL;
2739 }
2740 for (i=0 ; i < tuplesize ; i++) {
2741 Py_INCREF(Py_None);
2742 PyTuple_SET_ITEM(result, i, Py_None);
2743 }
2744
2745 /* create izipobject structure */
2746 lz = (izipobject *)type->tp_alloc(type, 0);
2747 if (lz == NULL) {
2748 Py_DECREF(ittuple);
2749 Py_DECREF(result);
2750 return NULL;
2751 }
2752 lz->ittuple = ittuple;
2753 lz->tuplesize = tuplesize;
2754 lz->result = result;
2755
2756 return (PyObject *)lz;
2757 }
2758
2759 static void
2760 izip_dealloc(izipobject *lz)
2761 {
2762 PyObject_GC_UnTrack(lz);
2763 Py_XDECREF(lz->ittuple);
2764 Py_XDECREF(lz->result);
2765 Py_TYPE(lz)->tp_free(lz);
2766 }
2767
2768 static int
2769 izip_traverse(izipobject *lz, visitproc visit, void *arg)
2770 {
2771 Py_VISIT(lz->ittuple);
2772 Py_VISIT(lz->result);
2773 return 0;
2774 }
2775
2776 static PyObject *
2777 izip_next(izipobject *lz)
2778 {
2779 Py_ssize_t i;
2780 Py_ssize_t tuplesize = lz->tuplesize;
2781 PyObject *result = lz->result;
2782 PyObject *it;
2783 PyObject *item;
2784 PyObject *olditem;
2785
2786 if (tuplesize == 0)
2787 return NULL;
2788 if (Py_REFCNT(result) == 1) {
2789 Py_INCREF(result);
2790 for (i=0 ; i < tuplesize ; i++) {
2791 it = PyTuple_GET_ITEM(lz->ittuple, i);
2792 assert(PyIter_Check(it));
2793 item = (*Py_TYPE(it)->tp_iternext)(it);
2794 if (item == NULL) {
2795 Py_DECREF(result);
2796 return NULL;
2797 }
2798 olditem = PyTuple_GET_ITEM(result, i);
2799 PyTuple_SET_ITEM(result, i, item);
2800 Py_DECREF(olditem);
2801 }
2802 } else {
2803 result = PyTuple_New(tuplesize);
2804 if (result == NULL)
2805 return NULL;
2806 for (i=0 ; i < tuplesize ; i++) {
2807 it = PyTuple_GET_ITEM(lz->ittuple, i);
2808 assert(PyIter_Check(it));
2809 item = (*Py_TYPE(it)->tp_iternext)(it);
2810 if (item == NULL) {
2811 Py_DECREF(result);
2812 return NULL;
2813 }
2814 PyTuple_SET_ITEM(result, i, item);
2815 }
2816 }
2817 return result;
2818 }
2819
2820 PyDoc_STRVAR(izip_doc,
2821 "izip(iter1 [,iter2 [...]]) --> izip object\n\
2822 \n\
2823 Return a izip object whose .next() method returns a tuple where\n\
2824 the i-th element comes from the i-th iterable argument. The .next()\n\
2825 method continues until the shortest iterable in the argument sequence\n\
2826 is exhausted and then it raises StopIteration. Works like the zip()\n\
2827 function but consumes less memory by returning an iterator instead of\n\
2828 a list.");
2829
2830 static PyTypeObject izip_type = {
2831 PyVarObject_HEAD_INIT(NULL, 0)
2832 "itertools.izip", /* tp_name */
2833 sizeof(izipobject), /* tp_basicsize */
2834 0, /* tp_itemsize */
2835 /* methods */
2836 (destructor)izip_dealloc, /* tp_dealloc */
2837 0, /* tp_print */
2838 0, /* tp_getattr */
2839 0, /* tp_setattr */
2840 0, /* tp_compare */
2841 0, /* tp_repr */
2842 0, /* tp_as_number */
2843 0, /* tp_as_sequence */
2844 0, /* tp_as_mapping */
2845 0, /* tp_hash */
2846 0, /* tp_call */
2847 0, /* tp_str */
2848 PyObject_GenericGetAttr, /* tp_getattro */
2849 0, /* tp_setattro */
2850 0, /* tp_as_buffer */
2851 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2852 Py_TPFLAGS_BASETYPE, /* tp_flags */
2853 izip_doc, /* tp_doc */
2854 (traverseproc)izip_traverse, /* tp_traverse */
2855 0, /* tp_clear */
2856 0, /* tp_richcompare */
2857 0, /* tp_weaklistoffset */
2858 PyObject_SelfIter, /* tp_iter */
2859 (iternextfunc)izip_next, /* tp_iternext */
2860 0, /* tp_methods */
2861 0, /* tp_members */
2862 0, /* tp_getset */
2863 0, /* tp_base */
2864 0, /* tp_dict */
2865 0, /* tp_descr_get */
2866 0, /* tp_descr_set */
2867 0, /* tp_dictoffset */
2868 0, /* tp_init */
2869 0, /* tp_alloc */
2870 izip_new, /* tp_new */
2871 PyObject_GC_Del, /* tp_free */
2872 };
2873
2874
2875 /* repeat object ************************************************************/
2876
2877 typedef struct {
2878 PyObject_HEAD
2879 PyObject *element;
2880 Py_ssize_t cnt;
2881 } repeatobject;
2882
2883 static PyTypeObject repeat_type;
2884
2885 static PyObject *
2886 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2887 {
2888 repeatobject *ro;
2889 PyObject *element;
2890 Py_ssize_t cnt = -1;
2891
2892 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
2893 return NULL;
2894
2895 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
2896 return NULL;
2897
2898 if (PyTuple_Size(args) == 2 && cnt < 0)
2899 cnt = 0;
2900
2901 ro = (repeatobject *)type->tp_alloc(type, 0);
2902 if (ro == NULL)
2903 return NULL;
2904 Py_INCREF(element);
2905 ro->element = element;
2906 ro->cnt = cnt;
2907 return (PyObject *)ro;
2908 }
2909
2910 static void
2911 repeat_dealloc(repeatobject *ro)
2912 {
2913 PyObject_GC_UnTrack(ro);
2914 Py_XDECREF(ro->element);
2915 Py_TYPE(ro)->tp_free(ro);
2916 }
2917
2918 static int
2919 repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2920 {
2921 Py_VISIT(ro->element);
2922 return 0;
2923 }
2924
2925 static PyObject *
2926 repeat_next(repeatobject *ro)
2927 {
2928 if (ro->cnt == 0)
2929 return NULL;
2930 if (ro->cnt > 0)
2931 ro->cnt--;
2932 Py_INCREF(ro->element);
2933 return ro->element;
2934 }
2935
2936 static PyObject *
2937 repeat_repr(repeatobject *ro)
2938 {
2939 PyObject *result, *objrepr;
2940
2941 objrepr = PyObject_Repr(ro->element);
2942 if (objrepr == NULL)
2943 return NULL;
2944
2945 if (ro->cnt == -1)
2946 result = PyString_FromFormat("repeat(%s)",
2947 PyString_AS_STRING(objrepr));
2948 else
2949 result = PyString_FromFormat("repeat(%s, %zd)",
2950 PyString_AS_STRING(objrepr), ro->cnt);
2951 Py_DECREF(objrepr);
2952 return result;
2953 }
2954
2955 static PyObject *
2956 repeat_len(repeatobject *ro)
2957 {
2958 if (ro->cnt == -1) {
2959 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2960 return NULL;
2961 }
2962 return PyInt_FromSize_t(ro->cnt);
2963 }
2964
2965 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
2966
2967 static PyMethodDef repeat_methods[] = {
2968 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
2969 {NULL, NULL} /* sentinel */
2970 };
2971
2972 PyDoc_STRVAR(repeat_doc,
2973 "repeat(element [,times]) -> create an iterator which returns the element\n\
2974 for the specified number of times. If not specified, returns the element\n\
2975 endlessly.");
2976
2977 static PyTypeObject repeat_type = {
2978 PyVarObject_HEAD_INIT(NULL, 0)
2979 "itertools.repeat", /* tp_name */
2980 sizeof(repeatobject), /* tp_basicsize */
2981 0, /* tp_itemsize */
2982 /* methods */
2983 (destructor)repeat_dealloc, /* tp_dealloc */
2984 0, /* tp_print */
2985 0, /* tp_getattr */
2986 0, /* tp_setattr */
2987 0, /* tp_compare */
2988 (reprfunc)repeat_repr, /* tp_repr */
2989 0, /* tp_as_number */
2990 0, /* tp_as_sequence */
2991 0, /* tp_as_mapping */
2992 0, /* tp_hash */
2993 0, /* tp_call */
2994 0, /* tp_str */
2995 PyObject_GenericGetAttr, /* tp_getattro */
2996 0, /* tp_setattro */
2997 0, /* tp_as_buffer */
2998 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2999 Py_TPFLAGS_BASETYPE, /* tp_flags */
3000 repeat_doc, /* tp_doc */
3001 (traverseproc)repeat_traverse, /* tp_traverse */
3002 0, /* tp_clear */
3003 0, /* tp_richcompare */
3004 0, /* tp_weaklistoffset */
3005 PyObject_SelfIter, /* tp_iter */
3006 (iternextfunc)repeat_next, /* tp_iternext */
3007 repeat_methods, /* tp_methods */
3008 0, /* tp_members */
3009 0, /* tp_getset */
3010 0, /* tp_base */
3011 0, /* tp_dict */
3012 0, /* tp_descr_get */
3013 0, /* tp_descr_set */
3014 0, /* tp_dictoffset */
3015 0, /* tp_init */
3016 0, /* tp_alloc */
3017 repeat_new, /* tp_new */
3018 PyObject_GC_Del, /* tp_free */
3019 };
3020
3021 /* iziplongest object ************************************************************/
3022
3023 #include "Python.h"
3024
3025 typedef struct {
3026 PyObject_HEAD
3027 Py_ssize_t tuplesize;
3028 Py_ssize_t numactive;
3029 PyObject *ittuple; /* tuple of iterators */
3030 PyObject *result;
3031 PyObject *fillvalue;
3032 } iziplongestobject;
3033
3034 static PyTypeObject iziplongest_type;
3035
3036 static PyObject *
3037 izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3038 {
3039 iziplongestobject *lz;
3040 Py_ssize_t i;
3041 PyObject *ittuple; /* tuple of iterators */
3042 PyObject *result;
3043 PyObject *fillvalue = Py_None;
3044 Py_ssize_t tuplesize = PySequence_Length(args);
3045
3046 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3047 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3048 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3049 PyErr_SetString(PyExc_TypeError,
3050 "izip_longest() got an unexpected keyword argument");
3051 return NULL;
3052 }
3053 }
3054
3055 /* args must be a tuple */
3056 assert(PyTuple_Check(args));
3057
3058 /* obtain iterators */
3059 ittuple = PyTuple_New(tuplesize);
3060 if (ittuple == NULL)
3061 return NULL;
3062 for (i=0; i < tuplesize; ++i) {
3063 PyObject *item = PyTuple_GET_ITEM(args, i);
3064 PyObject *it = PyObject_GetIter(item);
3065 if (it == NULL) {
3066 if (PyErr_ExceptionMatches(PyExc_TypeError))
3067 PyErr_Format(PyExc_TypeError,
3068 "izip_longest argument #%zd must support iteration",
3069 i+1);
3070 Py_DECREF(ittuple);
3071 return NULL;
3072 }
3073 PyTuple_SET_ITEM(ittuple, i, it);
3074 }
3075
3076 /* create a result holder */
3077 result = PyTuple_New(tuplesize);
3078 if (result == NULL) {
3079 Py_DECREF(ittuple);
3080 return NULL;
3081 }
3082 for (i=0 ; i < tuplesize ; i++) {
3083 Py_INCREF(Py_None);
3084 PyTuple_SET_ITEM(result, i, Py_None);
3085 }
3086
3087 /* create iziplongestobject structure */
3088 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3089 if (lz == NULL) {
3090 Py_DECREF(ittuple);
3091 Py_DECREF(result);
3092 return NULL;
3093 }
3094 lz->ittuple = ittuple;
3095 lz->tuplesize = tuplesize;
3096 lz->numactive = tuplesize;
3097 lz->result = result;
3098 Py_INCREF(fillvalue);
3099 lz->fillvalue = fillvalue;
3100 return (PyObject *)lz;
3101 }
3102
3103 static void
3104 izip_longest_dealloc(iziplongestobject *lz)
3105 {
3106 PyObject_GC_UnTrack(lz);
3107 Py_XDECREF(lz->ittuple);
3108 Py_XDECREF(lz->result);
3109 Py_XDECREF(lz->fillvalue);
3110 Py_TYPE(lz)->tp_free(lz);
3111 }
3112
3113 static int
3114 izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3115 {
3116 Py_VISIT(lz->ittuple);
3117 Py_VISIT(lz->result);
3118 Py_VISIT(lz->fillvalue);
3119 return 0;
3120 }
3121
3122 static PyObject *
3123 izip_longest_next(iziplongestobject *lz)
3124 {
3125 Py_ssize_t i;
3126 Py_ssize_t tuplesize = lz->tuplesize;
3127 PyObject *result = lz->result;
3128 PyObject *it;
3129 PyObject *item;
3130 PyObject *olditem;
3131
3132 if (tuplesize == 0)
3133 return NULL;
3134 if (lz->numactive == 0)
3135 return NULL;
3136 if (Py_REFCNT(result) == 1) {
3137 Py_INCREF(result);
3138 for (i=0 ; i < tuplesize ; i++) {
3139 it = PyTuple_GET_ITEM(lz->ittuple, i);
3140 if (it == NULL) {
3141 Py_INCREF(lz->fillvalue);
3142 item = lz->fillvalue;
3143 } else {
3144 assert(PyIter_Check(it));
3145 item = (*Py_TYPE(it)->tp_iternext)(it);
3146 if (item == NULL) {
3147 lz->numactive -= 1;
3148 if (lz->numactive == 0) {
3149 Py_DECREF(result);
3150 return NULL;
3151 } else {
3152 Py_INCREF(lz->fillvalue);
3153 item = lz->fillvalue;
3154 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3155 Py_DECREF(it);
3156 }
3157 }
3158 }
3159 olditem = PyTuple_GET_ITEM(result, i);
3160 PyTuple_SET_ITEM(result, i, item);
3161 Py_DECREF(olditem);
3162 }
3163 } else {
3164 result = PyTuple_New(tuplesize);
3165 if (result == NULL)
3166 return NULL;
3167 for (i=0 ; i < tuplesize ; i++) {
3168 it = PyTuple_GET_ITEM(lz->ittuple, i);
3169 if (it == NULL) {
3170 Py_INCREF(lz->fillvalue);
3171 item = lz->fillvalue;
3172 } else {
3173 assert(PyIter_Check(it));
3174 item = (*Py_TYPE(it)->tp_iternext)(it);
3175 if (item == NULL) {
3176 lz->numactive -= 1;
3177 if (lz->numactive == 0) {
3178 Py_DECREF(result);
3179 return NULL;
3180 } else {
3181 Py_INCREF(lz->fillvalue);
3182 item = lz->fillvalue;
3183 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3184 Py_DECREF(it);
3185 }
3186 }
3187 }
3188 PyTuple_SET_ITEM(result, i, item);
3189 }
3190 }
3191 return result;
3192 }
3193
3194 PyDoc_STRVAR(izip_longest_doc,
3195 "izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3196 \n\
3197 Return an izip_longest object whose .next() method returns a tuple where\n\
3198 the i-th element comes from the i-th iterable argument. The .next()\n\
3199 method continues until the longest iterable in the argument sequence\n\
3200 is exhausted and then it raises StopIteration. When the shorter iterables\n\
3201 are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3202 defaults to None or can be specified by a keyword argument.\n\
3203 ");
3204
3205 static PyTypeObject iziplongest_type = {
3206 PyVarObject_HEAD_INIT(NULL, 0)
3207 "itertools.izip_longest", /* tp_name */
3208 sizeof(iziplongestobject), /* tp_basicsize */
3209 0, /* tp_itemsize */
3210 /* methods */
3211 (destructor)izip_longest_dealloc, /* tp_dealloc */
3212 0, /* tp_print */
3213 0, /* tp_getattr */
3214 0, /* tp_setattr */
3215 0, /* tp_compare */
3216 0, /* tp_repr */
3217 0, /* tp_as_number */
3218 0, /* tp_as_sequence */
3219 0, /* tp_as_mapping */
3220 0, /* tp_hash */
3221 0, /* tp_call */
3222 0, /* tp_str */
3223 PyObject_GenericGetAttr, /* tp_getattro */
3224 0, /* tp_setattro */
3225 0, /* tp_as_buffer */
3226 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3227 Py_TPFLAGS_BASETYPE, /* tp_flags */
3228 izip_longest_doc, /* tp_doc */
3229 (traverseproc)izip_longest_traverse, /* tp_traverse */
3230 0, /* tp_clear */
3231 0, /* tp_richcompare */
3232 0, /* tp_weaklistoffset */
3233 PyObject_SelfIter, /* tp_iter */
3234 (iternextfunc)izip_longest_next, /* tp_iternext */
3235 0, /* tp_methods */
3236 0, /* tp_members */
3237 0, /* tp_getset */
3238 0, /* tp_base */
3239 0, /* tp_dict */
3240 0, /* tp_descr_get */
3241 0, /* tp_descr_set */
3242 0, /* tp_dictoffset */
3243 0, /* tp_init */
3244 0, /* tp_alloc */
3245 izip_longest_new, /* tp_new */
3246 PyObject_GC_Del, /* tp_free */
3247 };
3248
3249 /* module level code ********************************************************/
3250
3251 PyDoc_STRVAR(module_doc,
3252 "Functional tools for creating and using iterators.\n\
3253 \n\
3254 Infinite iterators:\n\
3255 count([n]) --> n, n+1, n+2, ...\n\
3256 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
3257 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
3258 \n\
3259 Iterators terminating on the shortest input sequence:\n\
3260 izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3261 izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3262 ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3263 ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
3264 islice(seq, [start,] stop [, step]) --> elements from\n\
3265 seq[start:stop:step]\n\
3266 imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3267 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
3268 tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
3269 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3270 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3271 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3272 groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
3273 ");
3274
3275
3276 static PyMethodDef module_methods[] = {
3277 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3278 {NULL, NULL} /* sentinel */
3279 };
3280
3281 PyMODINIT_FUNC
3282 inititertools(void)
3283 {
3284 int i;
3285 PyObject *m;
3286 char *name;
3287 PyTypeObject *typelist[] = {
3288 &combinations_type,
3289 &cycle_type,
3290 &dropwhile_type,
3291 &takewhile_type,
3292 &islice_type,
3293 &starmap_type,
3294 &imap_type,
3295 &chain_type,
3296 &ifilter_type,
3297 &ifilterfalse_type,
3298 &count_type,
3299 &izip_type,
3300 &iziplongest_type,
3301 &product_type,
3302 &repeat_type,
3303 &groupby_type,
3304 NULL
3305 };
3306
3307 Py_TYPE(&teedataobject_type) = &PyType_Type;
3308 m = Py_InitModule3("itertools", module_methods, module_doc);
3309 if (m == NULL)
3310 return;
3311
3312 for (i=0 ; typelist[i] != NULL ; i++) {
3313 if (PyType_Ready(typelist[i]) < 0)
3314 return;
3315 name = strchr(typelist[i]->tp_name, '.');
3316 assert (name != NULL);
3317 Py_INCREF(typelist[i]);
3318 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3319 }
3320
3321 if (PyType_Ready(&teedataobject_type) < 0)
3322 return;
3323 if (PyType_Ready(&tee_type) < 0)
3324 return;
3325 if (PyType_Ready(&_grouper_type) < 0)
3326 return;
3327 }
python git mirror