No issues found
1 #include "Python.h"
2 #include "structmember.h"
3
4 /* Itertools module written and maintained
5 by Raymond D. Hettinger <python@rcn.com>
6 Copyright (c) 2003 Python Software Foundation.
7 All rights reserved.
8 */
9
10
11 /* groupby object ***********************************************************/
12
13 typedef struct {
14 PyObject_HEAD
15 PyObject *it;
16 PyObject *keyfunc;
17 PyObject *tgtkey;
18 PyObject *currkey;
19 PyObject *currvalue;
20 } groupbyobject;
21
22 static PyTypeObject groupby_type;
23 static PyObject *_grouper_create(groupbyobject *, PyObject *);
24
25 static PyObject *
26 groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
27 {
28 static char *kwargs[] = {"iterable", "key", NULL};
29 groupbyobject *gbo;
30 PyObject *it, *keyfunc = Py_None;
31
32 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
33 &it, &keyfunc))
34 return NULL;
35
36 gbo = (groupbyobject *)type->tp_alloc(type, 0);
37 if (gbo == NULL)
38 return NULL;
39 gbo->tgtkey = NULL;
40 gbo->currkey = NULL;
41 gbo->currvalue = NULL;
42 gbo->keyfunc = keyfunc;
43 Py_INCREF(keyfunc);
44 gbo->it = PyObject_GetIter(it);
45 if (gbo->it == NULL) {
46 Py_DECREF(gbo);
47 return NULL;
48 }
49 return (PyObject *)gbo;
50 }
51
52 static void
53 groupby_dealloc(groupbyobject *gbo)
54 {
55 PyObject_GC_UnTrack(gbo);
56 Py_XDECREF(gbo->it);
57 Py_XDECREF(gbo->keyfunc);
58 Py_XDECREF(gbo->tgtkey);
59 Py_XDECREF(gbo->currkey);
60 Py_XDECREF(gbo->currvalue);
61 Py_TYPE(gbo)->tp_free(gbo);
62 }
63
64 static int
65 groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
66 {
67 Py_VISIT(gbo->it);
68 Py_VISIT(gbo->keyfunc);
69 Py_VISIT(gbo->tgtkey);
70 Py_VISIT(gbo->currkey);
71 Py_VISIT(gbo->currvalue);
72 return 0;
73 }
74
75 static PyObject *
76 groupby_next(groupbyobject *gbo)
77 {
78 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
79
80 /* skip to next iteration group */
81 for (;;) {
82 if (gbo->currkey == NULL)
83 /* pass */;
84 else if (gbo->tgtkey == NULL)
85 break;
86 else {
87 int rcmp;
88
89 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
90 gbo->currkey, Py_EQ);
91 if (rcmp == -1)
92 return NULL;
93 else if (rcmp == 0)
94 break;
95 }
96
97 newvalue = PyIter_Next(gbo->it);
98 if (newvalue == NULL)
99 return NULL;
100
101 if (gbo->keyfunc == Py_None) {
102 newkey = newvalue;
103 Py_INCREF(newvalue);
104 } else {
105 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
106 newvalue, NULL);
107 if (newkey == NULL) {
108 Py_DECREF(newvalue);
109 return NULL;
110 }
111 }
112
113 tmp = gbo->currkey;
114 gbo->currkey = newkey;
115 Py_XDECREF(tmp);
116
117 tmp = gbo->currvalue;
118 gbo->currvalue = newvalue;
119 Py_XDECREF(tmp);
120 }
121
122 Py_INCREF(gbo->currkey);
123 tmp = gbo->tgtkey;
124 gbo->tgtkey = gbo->currkey;
125 Py_XDECREF(tmp);
126
127 grouper = _grouper_create(gbo, gbo->tgtkey);
128 if (grouper == NULL)
129 return NULL;
130
131 r = PyTuple_Pack(2, gbo->currkey, grouper);
132 Py_DECREF(grouper);
133 return r;
134 }
135
136 PyDoc_STRVAR(groupby_doc,
137 "groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
138 (key, sub-iterator) grouped by each value of key(value).\n");
139
140 static PyTypeObject groupby_type = {
141 PyVarObject_HEAD_INIT(NULL, 0)
142 "itertools.groupby", /* tp_name */
143 sizeof(groupbyobject), /* tp_basicsize */
144 0, /* tp_itemsize */
145 /* methods */
146 (destructor)groupby_dealloc, /* tp_dealloc */
147 0, /* tp_print */
148 0, /* tp_getattr */
149 0, /* tp_setattr */
150 0, /* tp_compare */
151 0, /* tp_repr */
152 0, /* tp_as_number */
153 0, /* tp_as_sequence */
154 0, /* tp_as_mapping */
155 0, /* tp_hash */
156 0, /* tp_call */
157 0, /* tp_str */
158 PyObject_GenericGetAttr, /* tp_getattro */
159 0, /* tp_setattro */
160 0, /* tp_as_buffer */
161 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
162 Py_TPFLAGS_BASETYPE, /* tp_flags */
163 groupby_doc, /* tp_doc */
164 (traverseproc)groupby_traverse, /* tp_traverse */
165 0, /* tp_clear */
166 0, /* tp_richcompare */
167 0, /* tp_weaklistoffset */
168 PyObject_SelfIter, /* tp_iter */
169 (iternextfunc)groupby_next, /* tp_iternext */
170 0, /* tp_methods */
171 0, /* tp_members */
172 0, /* tp_getset */
173 0, /* tp_base */
174 0, /* tp_dict */
175 0, /* tp_descr_get */
176 0, /* tp_descr_set */
177 0, /* tp_dictoffset */
178 0, /* tp_init */
179 0, /* tp_alloc */
180 groupby_new, /* tp_new */
181 PyObject_GC_Del, /* tp_free */
182 };
183
184
185 /* _grouper object (internal) ************************************************/
186
187 typedef struct {
188 PyObject_HEAD
189 PyObject *parent;
190 PyObject *tgtkey;
191 } _grouperobject;
192
193 static PyTypeObject _grouper_type;
194
195 static PyObject *
196 _grouper_create(groupbyobject *parent, PyObject *tgtkey)
197 {
198 _grouperobject *igo;
199
200 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
201 if (igo == NULL)
202 return NULL;
203 igo->parent = (PyObject *)parent;
204 Py_INCREF(parent);
205 igo->tgtkey = tgtkey;
206 Py_INCREF(tgtkey);
207
208 PyObject_GC_Track(igo);
209 return (PyObject *)igo;
210 }
211
212 static void
213 _grouper_dealloc(_grouperobject *igo)
214 {
215 PyObject_GC_UnTrack(igo);
216 Py_DECREF(igo->parent);
217 Py_DECREF(igo->tgtkey);
218 PyObject_GC_Del(igo);
219 }
220
221 static int
222 _grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
223 {
224 Py_VISIT(igo->parent);
225 Py_VISIT(igo->tgtkey);
226 return 0;
227 }
228
229 static PyObject *
230 _grouper_next(_grouperobject *igo)
231 {
232 groupbyobject *gbo = (groupbyobject *)igo->parent;
233 PyObject *newvalue, *newkey, *r;
234 int rcmp;
235
236 if (gbo->currvalue == NULL) {
237 newvalue = PyIter_Next(gbo->it);
238 if (newvalue == NULL)
239 return NULL;
240
241 if (gbo->keyfunc == Py_None) {
242 newkey = newvalue;
243 Py_INCREF(newvalue);
244 } else {
245 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
246 newvalue, NULL);
247 if (newkey == NULL) {
248 Py_DECREF(newvalue);
249 return NULL;
250 }
251 }
252
253 assert(gbo->currkey == NULL);
254 gbo->currkey = newkey;
255 gbo->currvalue = newvalue;
256 }
257
258 assert(gbo->currkey != NULL);
259 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
260 if (rcmp <= 0)
261 /* got any error or current group is end */
262 return NULL;
263
264 r = gbo->currvalue;
265 gbo->currvalue = NULL;
266 Py_CLEAR(gbo->currkey);
267
268 return r;
269 }
270
271 static PyTypeObject _grouper_type = {
272 PyVarObject_HEAD_INIT(NULL, 0)
273 "itertools._grouper", /* tp_name */
274 sizeof(_grouperobject), /* tp_basicsize */
275 0, /* tp_itemsize */
276 /* methods */
277 (destructor)_grouper_dealloc, /* tp_dealloc */
278 0, /* tp_print */
279 0, /* tp_getattr */
280 0, /* tp_setattr */
281 0, /* tp_compare */
282 0, /* tp_repr */
283 0, /* tp_as_number */
284 0, /* tp_as_sequence */
285 0, /* tp_as_mapping */
286 0, /* tp_hash */
287 0, /* tp_call */
288 0, /* tp_str */
289 PyObject_GenericGetAttr, /* tp_getattro */
290 0, /* tp_setattro */
291 0, /* tp_as_buffer */
292 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
293 0, /* tp_doc */
294 (traverseproc)_grouper_traverse,/* tp_traverse */
295 0, /* tp_clear */
296 0, /* tp_richcompare */
297 0, /* tp_weaklistoffset */
298 PyObject_SelfIter, /* tp_iter */
299 (iternextfunc)_grouper_next, /* tp_iternext */
300 0, /* tp_methods */
301 0, /* tp_members */
302 0, /* tp_getset */
303 0, /* tp_base */
304 0, /* tp_dict */
305 0, /* tp_descr_get */
306 0, /* tp_descr_set */
307 0, /* tp_dictoffset */
308 0, /* tp_init */
309 0, /* tp_alloc */
310 0, /* tp_new */
311 PyObject_GC_Del, /* tp_free */
312 };
313
314
315
316 /* tee object and with supporting function and objects ***************/
317
318 /* The teedataobject pre-allocates space for LINKCELLS number of objects.
319 To help the object fit neatly inside cache lines (space for 16 to 32
320 pointers), the value should be a multiple of 16 minus space for
321 the other structure members including PyHEAD overhead. The larger the
322 value, the less memory overhead per object and the less time spent
323 allocating/deallocating new links. The smaller the number, the less
324 wasted space and the more rapid freeing of older data.
325 */
326 #define LINKCELLS 57
327
328 typedef struct {
329 PyObject_HEAD
330 PyObject *it;
331 int numread;
332 PyObject *nextlink;
333 PyObject *(values[LINKCELLS]);
334 } teedataobject;
335
336 typedef struct {
337 PyObject_HEAD
338 teedataobject *dataobj;
339 int index;
340 PyObject *weakreflist;
341 } teeobject;
342
343 static PyTypeObject teedataobject_type;
344
345 static PyObject *
346 teedataobject_new(PyObject *it)
347 {
348 teedataobject *tdo;
349
350 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
351 if (tdo == NULL)
352 return NULL;
353
354 tdo->numread = 0;
355 tdo->nextlink = NULL;
356 Py_INCREF(it);
357 tdo->it = it;
358 PyObject_GC_Track(tdo);
359 return (PyObject *)tdo;
360 }
361
362 static PyObject *
363 teedataobject_jumplink(teedataobject *tdo)
364 {
365 if (tdo->nextlink == NULL)
366 tdo->nextlink = teedataobject_new(tdo->it);
367 Py_XINCREF(tdo->nextlink);
368 return tdo->nextlink;
369 }
370
371 static PyObject *
372 teedataobject_getitem(teedataobject *tdo, int i)
373 {
374 PyObject *value;
375
376 assert(i < LINKCELLS);
377 if (i < tdo->numread)
378 value = tdo->values[i];
379 else {
380 /* this is the lead iterator, so fetch more data */
381 assert(i == tdo->numread);
382 value = PyIter_Next(tdo->it);
383 if (value == NULL)
384 return NULL;
385 tdo->numread++;
386 tdo->values[i] = value;
387 }
388 Py_INCREF(value);
389 return value;
390 }
391
392 static int
393 teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
394 {
395 int i;
396 Py_VISIT(tdo->it);
397 for (i = 0; i < tdo->numread; i++)
398 Py_VISIT(tdo->values[i]);
399 Py_VISIT(tdo->nextlink);
400 return 0;
401 }
402
403 static int
404 teedataobject_clear(teedataobject *tdo)
405 {
406 int i;
407 Py_CLEAR(tdo->it);
408 for (i=0 ; i<tdo->numread ; i++)
409 Py_CLEAR(tdo->values[i]);
410 Py_CLEAR(tdo->nextlink);
411 return 0;
412 }
413
414 static void
415 teedataobject_dealloc(teedataobject *tdo)
416 {
417 PyObject_GC_UnTrack(tdo);
418 teedataobject_clear(tdo);
419 PyObject_GC_Del(tdo);
420 }
421
422 PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
423
424 static PyTypeObject teedataobject_type = {
425 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
426 "itertools.tee_dataobject", /* tp_name */
427 sizeof(teedataobject), /* tp_basicsize */
428 0, /* tp_itemsize */
429 /* methods */
430 (destructor)teedataobject_dealloc, /* tp_dealloc */
431 0, /* tp_print */
432 0, /* tp_getattr */
433 0, /* tp_setattr */
434 0, /* tp_compare */
435 0, /* tp_repr */
436 0, /* tp_as_number */
437 0, /* tp_as_sequence */
438 0, /* tp_as_mapping */
439 0, /* tp_hash */
440 0, /* tp_call */
441 0, /* tp_str */
442 PyObject_GenericGetAttr, /* tp_getattro */
443 0, /* tp_setattro */
444 0, /* tp_as_buffer */
445 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
446 teedataobject_doc, /* tp_doc */
447 (traverseproc)teedataobject_traverse, /* tp_traverse */
448 (inquiry)teedataobject_clear, /* tp_clear */
449 0, /* tp_richcompare */
450 0, /* tp_weaklistoffset */
451 0, /* tp_iter */
452 0, /* tp_iternext */
453 0, /* tp_methods */
454 0, /* tp_members */
455 0, /* tp_getset */
456 0, /* tp_base */
457 0, /* tp_dict */
458 0, /* tp_descr_get */
459 0, /* tp_descr_set */
460 0, /* tp_dictoffset */
461 0, /* tp_init */
462 0, /* tp_alloc */
463 0, /* tp_new */
464 PyObject_GC_Del, /* tp_free */
465 };
466
467
468 static PyTypeObject tee_type;
469
470 static PyObject *
471 tee_next(teeobject *to)
472 {
473 PyObject *value, *link;
474
475 if (to->index >= LINKCELLS) {
476 link = teedataobject_jumplink(to->dataobj);
477 Py_DECREF(to->dataobj);
478 to->dataobj = (teedataobject *)link;
479 to->index = 0;
480 }
481 value = teedataobject_getitem(to->dataobj, to->index);
482 if (value == NULL)
483 return NULL;
484 to->index++;
485 return value;
486 }
487
488 static int
489 tee_traverse(teeobject *to, visitproc visit, void *arg)
490 {
491 Py_VISIT((PyObject *)to->dataobj);
492 return 0;
493 }
494
495 static PyObject *
496 tee_copy(teeobject *to)
497 {
498 teeobject *newto;
499
500 newto = PyObject_GC_New(teeobject, &tee_type);
501 if (newto == NULL)
502 return NULL;
503 Py_INCREF(to->dataobj);
504 newto->dataobj = to->dataobj;
505 newto->index = to->index;
506 newto->weakreflist = NULL;
507 PyObject_GC_Track(newto);
508 return (PyObject *)newto;
509 }
510
511 PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
512
513 static PyObject *
514 tee_fromiterable(PyObject *iterable)
515 {
516 teeobject *to;
517 PyObject *it = NULL;
518
519 it = PyObject_GetIter(iterable);
520 if (it == NULL)
521 return NULL;
522 if (PyObject_TypeCheck(it, &tee_type)) {
523 to = (teeobject *)tee_copy((teeobject *)it);
524 goto done;
525 }
526
527 to = PyObject_GC_New(teeobject, &tee_type);
528 if (to == NULL)
529 goto done;
530 to->dataobj = (teedataobject *)teedataobject_new(it);
531 if (!to->dataobj) {
532 PyObject_GC_Del(to);
533 to = NULL;
534 goto done;
535 }
536
537 to->index = 0;
538 to->weakreflist = NULL;
539 PyObject_GC_Track(to);
540 done:
541 Py_XDECREF(it);
542 return (PyObject *)to;
543 }
544
545 static PyObject *
546 tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
547 {
548 PyObject *iterable;
549
550 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
551 return NULL;
552 return tee_fromiterable(iterable);
553 }
554
555 static int
556 tee_clear(teeobject *to)
557 {
558 if (to->weakreflist != NULL)
559 PyObject_ClearWeakRefs((PyObject *) to);
560 Py_CLEAR(to->dataobj);
561 return 0;
562 }
563
564 static void
565 tee_dealloc(teeobject *to)
566 {
567 PyObject_GC_UnTrack(to);
568 tee_clear(to);
569 PyObject_GC_Del(to);
570 }
571
572 PyDoc_STRVAR(teeobject_doc,
573 "Iterator wrapped to make it copyable");
574
575 static PyMethodDef tee_methods[] = {
576 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
577 {NULL, NULL} /* sentinel */
578 };
579
580 static PyTypeObject tee_type = {
581 PyVarObject_HEAD_INIT(NULL, 0)
582 "itertools.tee", /* tp_name */
583 sizeof(teeobject), /* tp_basicsize */
584 0, /* tp_itemsize */
585 /* methods */
586 (destructor)tee_dealloc, /* tp_dealloc */
587 0, /* tp_print */
588 0, /* tp_getattr */
589 0, /* tp_setattr */
590 0, /* tp_compare */
591 0, /* tp_repr */
592 0, /* tp_as_number */
593 0, /* tp_as_sequence */
594 0, /* tp_as_mapping */
595 0, /* tp_hash */
596 0, /* tp_call */
597 0, /* tp_str */
598 0, /* tp_getattro */
599 0, /* tp_setattro */
600 0, /* tp_as_buffer */
601 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
602 teeobject_doc, /* tp_doc */
603 (traverseproc)tee_traverse, /* tp_traverse */
604 (inquiry)tee_clear, /* tp_clear */
605 0, /* tp_richcompare */
606 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
607 PyObject_SelfIter, /* tp_iter */
608 (iternextfunc)tee_next, /* tp_iternext */
609 tee_methods, /* tp_methods */
610 0, /* tp_members */
611 0, /* tp_getset */
612 0, /* tp_base */
613 0, /* tp_dict */
614 0, /* tp_descr_get */
615 0, /* tp_descr_set */
616 0, /* tp_dictoffset */
617 0, /* tp_init */
618 0, /* tp_alloc */
619 tee_new, /* tp_new */
620 PyObject_GC_Del, /* tp_free */
621 };
622
623 static PyObject *
624 tee(PyObject *self, PyObject *args)
625 {
626 Py_ssize_t i, n=2;
627 PyObject *it, *iterable, *copyable, *result;
628
629 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
630 return NULL;
631 if (n < 0) {
632 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
633 return NULL;
634 }
635 result = PyTuple_New(n);
636 if (result == NULL)
637 return NULL;
638 if (n == 0)
639 return result;
640 it = PyObject_GetIter(iterable);
641 if (it == NULL) {
642 Py_DECREF(result);
643 return NULL;
644 }
645 if (!PyObject_HasAttrString(it, "__copy__")) {
646 copyable = tee_fromiterable(it);
647 Py_DECREF(it);
648 if (copyable == NULL) {
649 Py_DECREF(result);
650 return NULL;
651 }
652 } else
653 copyable = it;
654 PyTuple_SET_ITEM(result, 0, copyable);
655 for (i=1 ; i<n ; i++) {
656 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
657 if (copyable == NULL) {
658 Py_DECREF(result);
659 return NULL;
660 }
661 PyTuple_SET_ITEM(result, i, copyable);
662 }
663 return result;
664 }
665
666 PyDoc_STRVAR(tee_doc,
667 "tee(iterable, n=2) --> tuple of n independent iterators.");
668
669
670 /* cycle object **********************************************************/
671
672 typedef struct {
673 PyObject_HEAD
674 PyObject *it;
675 PyObject *saved;
676 int firstpass;
677 } cycleobject;
678
679 static PyTypeObject cycle_type;
680
681 static PyObject *
682 cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
683 {
684 PyObject *it;
685 PyObject *iterable;
686 PyObject *saved;
687 cycleobject *lz;
688
689 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
690 return NULL;
691
692 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
693 return NULL;
694
695 /* Get iterator. */
696 it = PyObject_GetIter(iterable);
697 if (it == NULL)
698 return NULL;
699
700 saved = PyList_New(0);
701 if (saved == NULL) {
702 Py_DECREF(it);
703 return NULL;
704 }
705
706 /* create cycleobject structure */
707 lz = (cycleobject *)type->tp_alloc(type, 0);
708 if (lz == NULL) {
709 Py_DECREF(it);
710 Py_DECREF(saved);
711 return NULL;
712 }
713 lz->it = it;
714 lz->saved = saved;
715 lz->firstpass = 0;
716
717 return (PyObject *)lz;
718 }
719
720 static void
721 cycle_dealloc(cycleobject *lz)
722 {
723 PyObject_GC_UnTrack(lz);
724 Py_XDECREF(lz->saved);
725 Py_XDECREF(lz->it);
726 Py_TYPE(lz)->tp_free(lz);
727 }
728
729 static int
730 cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
731 {
732 Py_VISIT(lz->it);
733 Py_VISIT(lz->saved);
734 return 0;
735 }
736
737 static PyObject *
738 cycle_next(cycleobject *lz)
739 {
740 PyObject *item;
741 PyObject *it;
742 PyObject *tmp;
743
744 while (1) {
745 item = PyIter_Next(lz->it);
746 if (item != NULL) {
747 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
748 Py_DECREF(item);
749 return NULL;
750 }
751 return item;
752 }
753 if (PyErr_Occurred()) {
754 if (PyErr_ExceptionMatches(PyExc_StopIteration))
755 PyErr_Clear();
756 else
757 return NULL;
758 }
759 if (PyList_Size(lz->saved) == 0)
760 return NULL;
761 it = PyObject_GetIter(lz->saved);
762 if (it == NULL)
763 return NULL;
764 tmp = lz->it;
765 lz->it = it;
766 lz->firstpass = 1;
767 Py_DECREF(tmp);
768 }
769 }
770
771 PyDoc_STRVAR(cycle_doc,
772 "cycle(iterable) --> cycle object\n\
773 \n\
774 Return elements from the iterable until it is exhausted.\n\
775 Then repeat the sequence indefinitely.");
776
777 static PyTypeObject cycle_type = {
778 PyVarObject_HEAD_INIT(NULL, 0)
779 "itertools.cycle", /* tp_name */
780 sizeof(cycleobject), /* tp_basicsize */
781 0, /* tp_itemsize */
782 /* methods */
783 (destructor)cycle_dealloc, /* tp_dealloc */
784 0, /* tp_print */
785 0, /* tp_getattr */
786 0, /* tp_setattr */
787 0, /* tp_compare */
788 0, /* tp_repr */
789 0, /* tp_as_number */
790 0, /* tp_as_sequence */
791 0, /* tp_as_mapping */
792 0, /* tp_hash */
793 0, /* tp_call */
794 0, /* tp_str */
795 PyObject_GenericGetAttr, /* tp_getattro */
796 0, /* tp_setattro */
797 0, /* tp_as_buffer */
798 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
799 Py_TPFLAGS_BASETYPE, /* tp_flags */
800 cycle_doc, /* tp_doc */
801 (traverseproc)cycle_traverse, /* tp_traverse */
802 0, /* tp_clear */
803 0, /* tp_richcompare */
804 0, /* tp_weaklistoffset */
805 PyObject_SelfIter, /* tp_iter */
806 (iternextfunc)cycle_next, /* tp_iternext */
807 0, /* tp_methods */
808 0, /* tp_members */
809 0, /* tp_getset */
810 0, /* tp_base */
811 0, /* tp_dict */
812 0, /* tp_descr_get */
813 0, /* tp_descr_set */
814 0, /* tp_dictoffset */
815 0, /* tp_init */
816 0, /* tp_alloc */
817 cycle_new, /* tp_new */
818 PyObject_GC_Del, /* tp_free */
819 };
820
821
822 /* dropwhile object **********************************************************/
823
824 typedef struct {
825 PyObject_HEAD
826 PyObject *func;
827 PyObject *it;
828 long start;
829 } dropwhileobject;
830
831 static PyTypeObject dropwhile_type;
832
833 static PyObject *
834 dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
835 {
836 PyObject *func, *seq;
837 PyObject *it;
838 dropwhileobject *lz;
839
840 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
841 return NULL;
842
843 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
844 return NULL;
845
846 /* Get iterator. */
847 it = PyObject_GetIter(seq);
848 if (it == NULL)
849 return NULL;
850
851 /* create dropwhileobject structure */
852 lz = (dropwhileobject *)type->tp_alloc(type, 0);
853 if (lz == NULL) {
854 Py_DECREF(it);
855 return NULL;
856 }
857 Py_INCREF(func);
858 lz->func = func;
859 lz->it = it;
860 lz->start = 0;
861
862 return (PyObject *)lz;
863 }
864
865 static void
866 dropwhile_dealloc(dropwhileobject *lz)
867 {
868 PyObject_GC_UnTrack(lz);
869 Py_XDECREF(lz->func);
870 Py_XDECREF(lz->it);
871 Py_TYPE(lz)->tp_free(lz);
872 }
873
874 static int
875 dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
876 {
877 Py_VISIT(lz->it);
878 Py_VISIT(lz->func);
879 return 0;
880 }
881
882 static PyObject *
883 dropwhile_next(dropwhileobject *lz)
884 {
885 PyObject *item, *good;
886 PyObject *it = lz->it;
887 long ok;
888 PyObject *(*iternext)(PyObject *);
889
890 iternext = *Py_TYPE(it)->tp_iternext;
891 for (;;) {
892 item = iternext(it);
893 if (item == NULL)
894 return NULL;
895 if (lz->start == 1)
896 return item;
897
898 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
899 if (good == NULL) {
900 Py_DECREF(item);
901 return NULL;
902 }
903 ok = PyObject_IsTrue(good);
904 Py_DECREF(good);
905 if (!ok) {
906 lz->start = 1;
907 return item;
908 }
909 Py_DECREF(item);
910 }
911 }
912
913 PyDoc_STRVAR(dropwhile_doc,
914 "dropwhile(predicate, iterable) --> dropwhile object\n\
915 \n\
916 Drop items from the iterable while predicate(item) is true.\n\
917 Afterwards, return every element until the iterable is exhausted.");
918
919 static PyTypeObject dropwhile_type = {
920 PyVarObject_HEAD_INIT(NULL, 0)
921 "itertools.dropwhile", /* tp_name */
922 sizeof(dropwhileobject), /* tp_basicsize */
923 0, /* tp_itemsize */
924 /* methods */
925 (destructor)dropwhile_dealloc, /* tp_dealloc */
926 0, /* tp_print */
927 0, /* tp_getattr */
928 0, /* tp_setattr */
929 0, /* tp_compare */
930 0, /* tp_repr */
931 0, /* tp_as_number */
932 0, /* tp_as_sequence */
933 0, /* tp_as_mapping */
934 0, /* tp_hash */
935 0, /* tp_call */
936 0, /* tp_str */
937 PyObject_GenericGetAttr, /* tp_getattro */
938 0, /* tp_setattro */
939 0, /* tp_as_buffer */
940 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
941 Py_TPFLAGS_BASETYPE, /* tp_flags */
942 dropwhile_doc, /* tp_doc */
943 (traverseproc)dropwhile_traverse, /* tp_traverse */
944 0, /* tp_clear */
945 0, /* tp_richcompare */
946 0, /* tp_weaklistoffset */
947 PyObject_SelfIter, /* tp_iter */
948 (iternextfunc)dropwhile_next, /* tp_iternext */
949 0, /* tp_methods */
950 0, /* tp_members */
951 0, /* tp_getset */
952 0, /* tp_base */
953 0, /* tp_dict */
954 0, /* tp_descr_get */
955 0, /* tp_descr_set */
956 0, /* tp_dictoffset */
957 0, /* tp_init */
958 0, /* tp_alloc */
959 dropwhile_new, /* tp_new */
960 PyObject_GC_Del, /* tp_free */
961 };
962
963
964 /* takewhile object **********************************************************/
965
966 typedef struct {
967 PyObject_HEAD
968 PyObject *func;
969 PyObject *it;
970 long stop;
971 } takewhileobject;
972
973 static PyTypeObject takewhile_type;
974
975 static PyObject *
976 takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
977 {
978 PyObject *func, *seq;
979 PyObject *it;
980 takewhileobject *lz;
981
982 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
983 return NULL;
984
985 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
986 return NULL;
987
988 /* Get iterator. */
989 it = PyObject_GetIter(seq);
990 if (it == NULL)
991 return NULL;
992
993 /* create takewhileobject structure */
994 lz = (takewhileobject *)type->tp_alloc(type, 0);
995 if (lz == NULL) {
996 Py_DECREF(it);
997 return NULL;
998 }
999 Py_INCREF(func);
1000 lz->func = func;
1001 lz->it = it;
1002 lz->stop = 0;
1003
1004 return (PyObject *)lz;
1005 }
1006
1007 static void
1008 takewhile_dealloc(takewhileobject *lz)
1009 {
1010 PyObject_GC_UnTrack(lz);
1011 Py_XDECREF(lz->func);
1012 Py_XDECREF(lz->it);
1013 Py_TYPE(lz)->tp_free(lz);
1014 }
1015
1016 static int
1017 takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1018 {
1019 Py_VISIT(lz->it);
1020 Py_VISIT(lz->func);
1021 return 0;
1022 }
1023
1024 static PyObject *
1025 takewhile_next(takewhileobject *lz)
1026 {
1027 PyObject *item, *good;
1028 PyObject *it = lz->it;
1029 long ok;
1030
1031 if (lz->stop == 1)
1032 return NULL;
1033
1034 item = (*Py_TYPE(it)->tp_iternext)(it);
1035 if (item == NULL)
1036 return NULL;
1037
1038 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1039 if (good == NULL) {
1040 Py_DECREF(item);
1041 return NULL;
1042 }
1043 ok = PyObject_IsTrue(good);
1044 Py_DECREF(good);
1045 if (ok)
1046 return item;
1047 Py_DECREF(item);
1048 lz->stop = 1;
1049 return NULL;
1050 }
1051
1052 PyDoc_STRVAR(takewhile_doc,
1053 "takewhile(predicate, iterable) --> takewhile object\n\
1054 \n\
1055 Return successive entries from an iterable as long as the \n\
1056 predicate evaluates to true for each entry.");
1057
1058 static PyTypeObject takewhile_type = {
1059 PyVarObject_HEAD_INIT(NULL, 0)
1060 "itertools.takewhile", /* tp_name */
1061 sizeof(takewhileobject), /* tp_basicsize */
1062 0, /* tp_itemsize */
1063 /* methods */
1064 (destructor)takewhile_dealloc, /* tp_dealloc */
1065 0, /* tp_print */
1066 0, /* tp_getattr */
1067 0, /* tp_setattr */
1068 0, /* tp_compare */
1069 0, /* tp_repr */
1070 0, /* tp_as_number */
1071 0, /* tp_as_sequence */
1072 0, /* tp_as_mapping */
1073 0, /* tp_hash */
1074 0, /* tp_call */
1075 0, /* tp_str */
1076 PyObject_GenericGetAttr, /* tp_getattro */
1077 0, /* tp_setattro */
1078 0, /* tp_as_buffer */
1079 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1080 Py_TPFLAGS_BASETYPE, /* tp_flags */
1081 takewhile_doc, /* tp_doc */
1082 (traverseproc)takewhile_traverse, /* tp_traverse */
1083 0, /* tp_clear */
1084 0, /* tp_richcompare */
1085 0, /* tp_weaklistoffset */
1086 PyObject_SelfIter, /* tp_iter */
1087 (iternextfunc)takewhile_next, /* tp_iternext */
1088 0, /* tp_methods */
1089 0, /* tp_members */
1090 0, /* tp_getset */
1091 0, /* tp_base */
1092 0, /* tp_dict */
1093 0, /* tp_descr_get */
1094 0, /* tp_descr_set */
1095 0, /* tp_dictoffset */
1096 0, /* tp_init */
1097 0, /* tp_alloc */
1098 takewhile_new, /* tp_new */
1099 PyObject_GC_Del, /* tp_free */
1100 };
1101
1102
1103 /* islice object ************************************************************/
1104
1105 typedef struct {
1106 PyObject_HEAD
1107 PyObject *it;
1108 Py_ssize_t next;
1109 Py_ssize_t stop;
1110 Py_ssize_t step;
1111 Py_ssize_t cnt;
1112 } isliceobject;
1113
1114 static PyTypeObject islice_type;
1115
1116 static PyObject *
1117 islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1118 {
1119 PyObject *seq;
1120 Py_ssize_t start=0, stop=-1, step=1;
1121 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1122 Py_ssize_t numargs;
1123 isliceobject *lz;
1124
1125 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1126 return NULL;
1127
1128 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1129 return NULL;
1130
1131 numargs = PyTuple_Size(args);
1132 if (numargs == 2) {
1133 if (a1 != Py_None) {
1134 stop = PyInt_AsSsize_t(a1);
1135 if (stop == -1) {
1136 if (PyErr_Occurred())
1137 PyErr_Clear();
1138 PyErr_SetString(PyExc_ValueError,
1139 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1140 return NULL;
1141 }
1142 }
1143 } else {
1144 if (a1 != Py_None)
1145 start = PyInt_AsSsize_t(a1);
1146 if (start == -1 && PyErr_Occurred())
1147 PyErr_Clear();
1148 if (a2 != Py_None) {
1149 stop = PyInt_AsSsize_t(a2);
1150 if (stop == -1) {
1151 if (PyErr_Occurred())
1152 PyErr_Clear();
1153 PyErr_SetString(PyExc_ValueError,
1154 "Stop argument for islice() must be None or an integer: 0 <= x <= maxint.");
1155 return NULL;
1156 }
1157 }
1158 }
1159 if (start<0 || stop<-1) {
1160 PyErr_SetString(PyExc_ValueError,
1161 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
1162 return NULL;
1163 }
1164
1165 if (a3 != NULL) {
1166 if (a3 != Py_None)
1167 step = PyInt_AsSsize_t(a3);
1168 if (step == -1 && PyErr_Occurred())
1169 PyErr_Clear();
1170 }
1171 if (step<1) {
1172 PyErr_SetString(PyExc_ValueError,
1173 "Step for islice() must be a positive integer or None.");
1174 return NULL;
1175 }
1176
1177 /* Get iterator. */
1178 it = PyObject_GetIter(seq);
1179 if (it == NULL)
1180 return NULL;
1181
1182 /* create isliceobject structure */
1183 lz = (isliceobject *)type->tp_alloc(type, 0);
1184 if (lz == NULL) {
1185 Py_DECREF(it);
1186 return NULL;
1187 }
1188 lz->it = it;
1189 lz->next = start;
1190 lz->stop = stop;
1191 lz->step = step;
1192 lz->cnt = 0L;
1193
1194 return (PyObject *)lz;
1195 }
1196
1197 static void
1198 islice_dealloc(isliceobject *lz)
1199 {
1200 PyObject_GC_UnTrack(lz);
1201 Py_XDECREF(lz->it);
1202 Py_TYPE(lz)->tp_free(lz);
1203 }
1204
1205 static int
1206 islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1207 {
1208 Py_VISIT(lz->it);
1209 return 0;
1210 }
1211
1212 static PyObject *
1213 islice_next(isliceobject *lz)
1214 {
1215 PyObject *item;
1216 PyObject *it = lz->it;
1217 Py_ssize_t stop = lz->stop;
1218 Py_ssize_t oldnext;
1219 PyObject *(*iternext)(PyObject *);
1220
1221 iternext = *Py_TYPE(it)->tp_iternext;
1222 while (lz->cnt < lz->next) {
1223 item = iternext(it);
1224 if (item == NULL)
1225 return NULL;
1226 Py_DECREF(item);
1227 lz->cnt++;
1228 }
1229 if (stop != -1 && lz->cnt >= stop)
1230 return NULL;
1231 item = iternext(it);
1232 if (item == NULL)
1233 return NULL;
1234 lz->cnt++;
1235 oldnext = lz->next;
1236 /* The (size_t) cast below avoids the danger of undefined
1237 behaviour from signed integer overflow. */
1238 lz->next += (size_t)lz->step;
1239 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1240 lz->next = stop;
1241 return item;
1242 }
1243
1244 PyDoc_STRVAR(islice_doc,
1245 "islice(iterable, [start,] stop [, step]) --> islice object\n\
1246 \n\
1247 Return an iterator whose next() method returns selected values from an\n\
1248 iterable. If start is specified, will skip all preceding elements;\n\
1249 otherwise, start defaults to zero. Step defaults to one. If\n\
1250 specified as another value, step determines how many values are \n\
1251 skipped between successive calls. Works like a slice() on a list\n\
1252 but returns an iterator.");
1253
1254 static PyTypeObject islice_type = {
1255 PyVarObject_HEAD_INIT(NULL, 0)
1256 "itertools.islice", /* tp_name */
1257 sizeof(isliceobject), /* tp_basicsize */
1258 0, /* tp_itemsize */
1259 /* methods */
1260 (destructor)islice_dealloc, /* tp_dealloc */
1261 0, /* tp_print */
1262 0, /* tp_getattr */
1263 0, /* tp_setattr */
1264 0, /* tp_compare */
1265 0, /* tp_repr */
1266 0, /* tp_as_number */
1267 0, /* tp_as_sequence */
1268 0, /* tp_as_mapping */
1269 0, /* tp_hash */
1270 0, /* tp_call */
1271 0, /* tp_str */
1272 PyObject_GenericGetAttr, /* tp_getattro */
1273 0, /* tp_setattro */
1274 0, /* tp_as_buffer */
1275 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1276 Py_TPFLAGS_BASETYPE, /* tp_flags */
1277 islice_doc, /* tp_doc */
1278 (traverseproc)islice_traverse, /* tp_traverse */
1279 0, /* tp_clear */
1280 0, /* tp_richcompare */
1281 0, /* tp_weaklistoffset */
1282 PyObject_SelfIter, /* tp_iter */
1283 (iternextfunc)islice_next, /* tp_iternext */
1284 0, /* tp_methods */
1285 0, /* tp_members */
1286 0, /* tp_getset */
1287 0, /* tp_base */
1288 0, /* tp_dict */
1289 0, /* tp_descr_get */
1290 0, /* tp_descr_set */
1291 0, /* tp_dictoffset */
1292 0, /* tp_init */
1293 0, /* tp_alloc */
1294 islice_new, /* tp_new */
1295 PyObject_GC_Del, /* tp_free */
1296 };
1297
1298
1299 /* starmap object ************************************************************/
1300
1301 typedef struct {
1302 PyObject_HEAD
1303 PyObject *func;
1304 PyObject *it;
1305 } starmapobject;
1306
1307 static PyTypeObject starmap_type;
1308
1309 static PyObject *
1310 starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1311 {
1312 PyObject *func, *seq;
1313 PyObject *it;
1314 starmapobject *lz;
1315
1316 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1317 return NULL;
1318
1319 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1320 return NULL;
1321
1322 /* Get iterator. */
1323 it = PyObject_GetIter(seq);
1324 if (it == NULL)
1325 return NULL;
1326
1327 /* create starmapobject structure */
1328 lz = (starmapobject *)type->tp_alloc(type, 0);
1329 if (lz == NULL) {
1330 Py_DECREF(it);
1331 return NULL;
1332 }
1333 Py_INCREF(func);
1334 lz->func = func;
1335 lz->it = it;
1336
1337 return (PyObject *)lz;
1338 }
1339
1340 static void
1341 starmap_dealloc(starmapobject *lz)
1342 {
1343 PyObject_GC_UnTrack(lz);
1344 Py_XDECREF(lz->func);
1345 Py_XDECREF(lz->it);
1346 Py_TYPE(lz)->tp_free(lz);
1347 }
1348
1349 static int
1350 starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1351 {
1352 Py_VISIT(lz->it);
1353 Py_VISIT(lz->func);
1354 return 0;
1355 }
1356
1357 static PyObject *
1358 starmap_next(starmapobject *lz)
1359 {
1360 PyObject *args;
1361 PyObject *result;
1362 PyObject *it = lz->it;
1363
1364 args = (*Py_TYPE(it)->tp_iternext)(it);
1365 if (args == NULL)
1366 return NULL;
1367 if (!PyTuple_CheckExact(args)) {
1368 PyObject *newargs = PySequence_Tuple(args);
1369 Py_DECREF(args);
1370 if (newargs == NULL)
1371 return NULL;
1372 args = newargs;
1373 }
1374 result = PyObject_Call(lz->func, args, NULL);
1375 Py_DECREF(args);
1376 return result;
1377 }
1378
1379 PyDoc_STRVAR(starmap_doc,
1380 "starmap(function, sequence) --> starmap object\n\
1381 \n\
1382 Return an iterator whose values are returned from the function evaluated\n\
1383 with a argument tuple taken from the given sequence.");
1384
1385 static PyTypeObject starmap_type = {
1386 PyVarObject_HEAD_INIT(NULL, 0)
1387 "itertools.starmap", /* tp_name */
1388 sizeof(starmapobject), /* tp_basicsize */
1389 0, /* tp_itemsize */
1390 /* methods */
1391 (destructor)starmap_dealloc, /* tp_dealloc */
1392 0, /* tp_print */
1393 0, /* tp_getattr */
1394 0, /* tp_setattr */
1395 0, /* tp_compare */
1396 0, /* tp_repr */
1397 0, /* tp_as_number */
1398 0, /* tp_as_sequence */
1399 0, /* tp_as_mapping */
1400 0, /* tp_hash */
1401 0, /* tp_call */
1402 0, /* tp_str */
1403 PyObject_GenericGetAttr, /* tp_getattro */
1404 0, /* tp_setattro */
1405 0, /* tp_as_buffer */
1406 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1407 Py_TPFLAGS_BASETYPE, /* tp_flags */
1408 starmap_doc, /* tp_doc */
1409 (traverseproc)starmap_traverse, /* tp_traverse */
1410 0, /* tp_clear */
1411 0, /* tp_richcompare */
1412 0, /* tp_weaklistoffset */
1413 PyObject_SelfIter, /* tp_iter */
1414 (iternextfunc)starmap_next, /* tp_iternext */
1415 0, /* tp_methods */
1416 0, /* tp_members */
1417 0, /* tp_getset */
1418 0, /* tp_base */
1419 0, /* tp_dict */
1420 0, /* tp_descr_get */
1421 0, /* tp_descr_set */
1422 0, /* tp_dictoffset */
1423 0, /* tp_init */
1424 0, /* tp_alloc */
1425 starmap_new, /* tp_new */
1426 PyObject_GC_Del, /* tp_free */
1427 };
1428
1429
1430 /* imap object ************************************************************/
1431
1432 typedef struct {
1433 PyObject_HEAD
1434 PyObject *iters;
1435 PyObject *func;
1436 } imapobject;
1437
1438 static PyTypeObject imap_type;
1439
1440 static PyObject *
1441 imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1442 {
1443 PyObject *it, *iters, *func;
1444 imapobject *lz;
1445 Py_ssize_t numargs, i;
1446
1447 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1448 return NULL;
1449
1450 numargs = PyTuple_Size(args);
1451 if (numargs < 2) {
1452 PyErr_SetString(PyExc_TypeError,
1453 "imap() must have at least two arguments.");
1454 return NULL;
1455 }
1456
1457 iters = PyTuple_New(numargs-1);
1458 if (iters == NULL)
1459 return NULL;
1460
1461 for (i=1 ; i<numargs ; i++) {
1462 /* Get iterator. */
1463 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1464 if (it == NULL) {
1465 Py_DECREF(iters);
1466 return NULL;
1467 }
1468 PyTuple_SET_ITEM(iters, i-1, it);
1469 }
1470
1471 /* create imapobject structure */
1472 lz = (imapobject *)type->tp_alloc(type, 0);
1473 if (lz == NULL) {
1474 Py_DECREF(iters);
1475 return NULL;
1476 }
1477 lz->iters = iters;
1478 func = PyTuple_GET_ITEM(args, 0);
1479 Py_INCREF(func);
1480 lz->func = func;
1481
1482 return (PyObject *)lz;
1483 }
1484
1485 static void
1486 imap_dealloc(imapobject *lz)
1487 {
1488 PyObject_GC_UnTrack(lz);
1489 Py_XDECREF(lz->iters);
1490 Py_XDECREF(lz->func);
1491 Py_TYPE(lz)->tp_free(lz);
1492 }
1493
1494 static int
1495 imap_traverse(imapobject *lz, visitproc visit, void *arg)
1496 {
1497 Py_VISIT(lz->iters);
1498 Py_VISIT(lz->func);
1499 return 0;
1500 }
1501
1502 /*
1503 imap() is an iterator version of __builtins__.map() except that it does
1504 not have the None fill-in feature. That was intentionally left out for
1505 the following reasons:
1506
1507 1) Itertools are designed to be easily combined and chained together.
1508 Having all tools stop with the shortest input is a unifying principle
1509 that makes it easier to combine finite iterators (supplying data) with
1510 infinite iterators like count() and repeat() (for supplying sequential
1511 or constant arguments to a function).
1512
1513 2) In typical use cases for combining itertools, having one finite data
1514 supplier run out before another is likely to be an error condition which
1515 should not pass silently by automatically supplying None.
1516
1517 3) The use cases for automatic None fill-in are rare -- not many functions
1518 do something useful when a parameter suddenly switches type and becomes
1519 None.
1520
1521 4) If a need does arise, it can be met by __builtins__.map() or by
1522 writing: chain(iterable, repeat(None)).
1523
1524 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1525 */
1526
1527 static PyObject *
1528 imap_next(imapobject *lz)
1529 {
1530 PyObject *val;
1531 PyObject *argtuple;
1532 PyObject *result;
1533 Py_ssize_t numargs, i;
1534
1535 numargs = PyTuple_Size(lz->iters);
1536 argtuple = PyTuple_New(numargs);
1537 if (argtuple == NULL)
1538 return NULL;
1539
1540 for (i=0 ; i<numargs ; i++) {
1541 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1542 if (val == NULL) {
1543 Py_DECREF(argtuple);
1544 return NULL;
1545 }
1546 PyTuple_SET_ITEM(argtuple, i, val);
1547 }
1548 if (lz->func == Py_None)
1549 return argtuple;
1550 result = PyObject_Call(lz->func, argtuple, NULL);
1551 Py_DECREF(argtuple);
1552 return result;
1553 }
1554
1555 PyDoc_STRVAR(imap_doc,
1556 "imap(func, *iterables) --> imap object\n\
1557 \n\
1558 Make an iterator that computes the function using arguments from\n\
1559 each of the iterables. Like map() except that it returns\n\
1560 an iterator instead of a list and that it stops when the shortest\n\
1561 iterable is exhausted instead of filling in None for shorter\n\
1562 iterables.");
1563
1564 static PyTypeObject imap_type = {
1565 PyVarObject_HEAD_INIT(NULL, 0)
1566 "itertools.imap", /* tp_name */
1567 sizeof(imapobject), /* tp_basicsize */
1568 0, /* tp_itemsize */
1569 /* methods */
1570 (destructor)imap_dealloc, /* tp_dealloc */
1571 0, /* tp_print */
1572 0, /* tp_getattr */
1573 0, /* tp_setattr */
1574 0, /* tp_compare */
1575 0, /* tp_repr */
1576 0, /* tp_as_number */
1577 0, /* tp_as_sequence */
1578 0, /* tp_as_mapping */
1579 0, /* tp_hash */
1580 0, /* tp_call */
1581 0, /* tp_str */
1582 PyObject_GenericGetAttr, /* tp_getattro */
1583 0, /* tp_setattro */
1584 0, /* tp_as_buffer */
1585 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1586 Py_TPFLAGS_BASETYPE, /* tp_flags */
1587 imap_doc, /* tp_doc */
1588 (traverseproc)imap_traverse, /* tp_traverse */
1589 0, /* tp_clear */
1590 0, /* tp_richcompare */
1591 0, /* tp_weaklistoffset */
1592 PyObject_SelfIter, /* tp_iter */
1593 (iternextfunc)imap_next, /* tp_iternext */
1594 0, /* tp_methods */
1595 0, /* tp_members */
1596 0, /* tp_getset */
1597 0, /* tp_base */
1598 0, /* tp_dict */
1599 0, /* tp_descr_get */
1600 0, /* tp_descr_set */
1601 0, /* tp_dictoffset */
1602 0, /* tp_init */
1603 0, /* tp_alloc */
1604 imap_new, /* tp_new */
1605 PyObject_GC_Del, /* tp_free */
1606 };
1607
1608
1609 /* chain object ************************************************************/
1610
1611 typedef struct {
1612 PyObject_HEAD
1613 PyObject *source; /* Iterator over input iterables */
1614 PyObject *active; /* Currently running input iterator */
1615 } chainobject;
1616
1617 static PyTypeObject chain_type;
1618
1619 static PyObject *
1620 chain_new_internal(PyTypeObject *type, PyObject *source)
1621 {
1622 chainobject *lz;
1623
1624 lz = (chainobject *)type->tp_alloc(type, 0);
1625 if (lz == NULL) {
1626 Py_DECREF(source);
1627 return NULL;
1628 }
1629
1630 lz->source = source;
1631 lz->active = NULL;
1632 return (PyObject *)lz;
1633 }
1634
1635 static PyObject *
1636 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1637 {
1638 PyObject *source;
1639
1640 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1641 return NULL;
1642
1643 source = PyObject_GetIter(args);
1644 if (source == NULL)
1645 return NULL;
1646
1647 return chain_new_internal(type, source);
1648 }
1649
1650 static PyObject *
1651 chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1652 {
1653 PyObject *source;
1654
1655 source = PyObject_GetIter(arg);
1656 if (source == NULL)
1657 return NULL;
1658
1659 return chain_new_internal(type, source);
1660 }
1661
1662 static void
1663 chain_dealloc(chainobject *lz)
1664 {
1665 PyObject_GC_UnTrack(lz);
1666 Py_XDECREF(lz->active);
1667 Py_XDECREF(lz->source);
1668 Py_TYPE(lz)->tp_free(lz);
1669 }
1670
1671 static int
1672 chain_traverse(chainobject *lz, visitproc visit, void *arg)
1673 {
1674 Py_VISIT(lz->source);
1675 Py_VISIT(lz->active);
1676 return 0;
1677 }
1678
1679 static PyObject *
1680 chain_next(chainobject *lz)
1681 {
1682 PyObject *item;
1683
1684 if (lz->source == NULL)
1685 return NULL; /* already stopped */
1686
1687 if (lz->active == NULL) {
1688 PyObject *iterable = PyIter_Next(lz->source);
1689 if (iterable == NULL) {
1690 Py_CLEAR(lz->source);
1691 return NULL; /* no more input sources */
1692 }
1693 lz->active = PyObject_GetIter(iterable);
1694 Py_DECREF(iterable);
1695 if (lz->active == NULL) {
1696 Py_CLEAR(lz->source);
1697 return NULL; /* input not iterable */
1698 }
1699 }
1700 item = PyIter_Next(lz->active);
1701 if (item != NULL)
1702 return item;
1703 if (PyErr_Occurred()) {
1704 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1705 PyErr_Clear();
1706 else
1707 return NULL; /* input raised an exception */
1708 }
1709 Py_CLEAR(lz->active);
1710 return chain_next(lz); /* recurse and use next active */
1711 }
1712
1713 PyDoc_STRVAR(chain_doc,
1714 "chain(*iterables) --> chain object\n\
1715 \n\
1716 Return a chain object whose .next() method returns elements from the\n\
1717 first iterable until it is exhausted, then elements from the next\n\
1718 iterable, until all of the iterables are exhausted.");
1719
1720 PyDoc_STRVAR(chain_from_iterable_doc,
1721 "chain.from_iterable(iterable) --> chain object\n\
1722 \n\
1723 Alternate chain() contructor taking a single iterable argument\n\
1724 that evaluates lazily.");
1725
1726 static PyMethodDef chain_methods[] = {
1727 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1728 chain_from_iterable_doc},
1729 {NULL, NULL} /* sentinel */
1730 };
1731
1732 static PyTypeObject chain_type = {
1733 PyVarObject_HEAD_INIT(NULL, 0)
1734 "itertools.chain", /* tp_name */
1735 sizeof(chainobject), /* tp_basicsize */
1736 0, /* tp_itemsize */
1737 /* methods */
1738 (destructor)chain_dealloc, /* tp_dealloc */
1739 0, /* tp_print */
1740 0, /* tp_getattr */
1741 0, /* tp_setattr */
1742 0, /* tp_compare */
1743 0, /* tp_repr */
1744 0, /* tp_as_number */
1745 0, /* tp_as_sequence */
1746 0, /* tp_as_mapping */
1747 0, /* tp_hash */
1748 0, /* tp_call */
1749 0, /* tp_str */
1750 PyObject_GenericGetAttr, /* tp_getattro */
1751 0, /* tp_setattro */
1752 0, /* tp_as_buffer */
1753 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1754 Py_TPFLAGS_BASETYPE, /* tp_flags */
1755 chain_doc, /* tp_doc */
1756 (traverseproc)chain_traverse, /* tp_traverse */
1757 0, /* tp_clear */
1758 0, /* tp_richcompare */
1759 0, /* tp_weaklistoffset */
1760 PyObject_SelfIter, /* tp_iter */
1761 (iternextfunc)chain_next, /* tp_iternext */
1762 chain_methods, /* tp_methods */
1763 0, /* tp_members */
1764 0, /* tp_getset */
1765 0, /* tp_base */
1766 0, /* tp_dict */
1767 0, /* tp_descr_get */
1768 0, /* tp_descr_set */
1769 0, /* tp_dictoffset */
1770 0, /* tp_init */
1771 0, /* tp_alloc */
1772 chain_new, /* tp_new */
1773 PyObject_GC_Del, /* tp_free */
1774 };
1775
1776
1777 /* product object ************************************************************/
1778
1779 typedef struct {
1780 PyObject_HEAD
1781 PyObject *pools; /* tuple of pool tuples */
1782 Py_ssize_t *indices; /* one index per pool */
1783 PyObject *result; /* most recently returned result tuple */
1784 int stopped; /* set to 1 when the product iterator is exhausted */
1785 } productobject;
1786
1787 static PyTypeObject product_type;
1788
1789 static PyObject *
1790 product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1791 {
1792 productobject *lz;
1793 Py_ssize_t nargs, npools, repeat=1;
1794 PyObject *pools = NULL;
1795 Py_ssize_t *indices = NULL;
1796 Py_ssize_t i;
1797
1798 if (kwds != NULL) {
1799 char *kwlist[] = {"repeat", 0};
1800 PyObject *tmpargs = PyTuple_New(0);
1801 if (tmpargs == NULL)
1802 return NULL;
1803 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1804 Py_DECREF(tmpargs);
1805 return NULL;
1806 }
1807 Py_DECREF(tmpargs);
1808 if (repeat < 0) {
1809 PyErr_SetString(PyExc_ValueError,
1810 "repeat argument cannot be negative");
1811 return NULL;
1812 }
1813 }
1814
1815 assert(PyTuple_Check(args));
1816 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1817 npools = nargs * repeat;
1818
1819 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1820 if (indices == NULL) {
1821 PyErr_NoMemory();
1822 goto error;
1823 }
1824
1825 pools = PyTuple_New(npools);
1826 if (pools == NULL)
1827 goto error;
1828
1829 for (i=0; i < nargs ; ++i) {
1830 PyObject *item = PyTuple_GET_ITEM(args, i);
1831 PyObject *pool = PySequence_Tuple(item);
1832 if (pool == NULL)
1833 goto error;
1834 PyTuple_SET_ITEM(pools, i, pool);
1835 indices[i] = 0;
1836 }
1837 for ( ; i < npools; ++i) {
1838 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1839 Py_INCREF(pool);
1840 PyTuple_SET_ITEM(pools, i, pool);
1841 indices[i] = 0;
1842 }
1843
1844 /* create productobject structure */
1845 lz = (productobject *)type->tp_alloc(type, 0);
1846 if (lz == NULL)
1847 goto error;
1848
1849 lz->pools = pools;
1850 lz->indices = indices;
1851 lz->result = NULL;
1852 lz->stopped = 0;
1853
1854 return (PyObject *)lz;
1855
1856 error:
1857 if (indices != NULL)
1858 PyMem_Free(indices);
1859 Py_XDECREF(pools);
1860 return NULL;
1861 }
1862
1863 static void
1864 product_dealloc(productobject *lz)
1865 {
1866 PyObject_GC_UnTrack(lz);
1867 Py_XDECREF(lz->pools);
1868 Py_XDECREF(lz->result);
1869 if (lz->indices != NULL)
1870 PyMem_Free(lz->indices);
1871 Py_TYPE(lz)->tp_free(lz);
1872 }
1873
1874 static int
1875 product_traverse(productobject *lz, visitproc visit, void *arg)
1876 {
1877 Py_VISIT(lz->pools);
1878 Py_VISIT(lz->result);
1879 return 0;
1880 }
1881
1882 static PyObject *
1883 product_next(productobject *lz)
1884 {
1885 PyObject *pool;
1886 PyObject *elem;
1887 PyObject *oldelem;
1888 PyObject *pools = lz->pools;
1889 PyObject *result = lz->result;
1890 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1891 Py_ssize_t i;
1892
1893 if (lz->stopped)
1894 return NULL;
1895
1896 if (result == NULL) {
1897 /* On the first pass, return an initial tuple filled with the
1898 first element from each pool. */
1899 result = PyTuple_New(npools);
1900 if (result == NULL)
1901 goto empty;
1902 lz->result = result;
1903 for (i=0; i < npools; i++) {
1904 pool = PyTuple_GET_ITEM(pools, i);
1905 if (PyTuple_GET_SIZE(pool) == 0)
1906 goto empty;
1907 elem = PyTuple_GET_ITEM(pool, 0);
1908 Py_INCREF(elem);
1909 PyTuple_SET_ITEM(result, i, elem);
1910 }
1911 } else {
1912 Py_ssize_t *indices = lz->indices;
1913
1914 /* Copy the previous result tuple or re-use it if available */
1915 if (Py_REFCNT(result) > 1) {
1916 PyObject *old_result = result;
1917 result = PyTuple_New(npools);
1918 if (result == NULL)
1919 goto empty;
1920 lz->result = result;
1921 for (i=0; i < npools; i++) {
1922 elem = PyTuple_GET_ITEM(old_result, i);
1923 Py_INCREF(elem);
1924 PyTuple_SET_ITEM(result, i, elem);
1925 }
1926 Py_DECREF(old_result);
1927 }
1928 /* Now, we've got the only copy so we can update it in-place */
1929 assert (npools==0 || Py_REFCNT(result) == 1);
1930
1931 /* Update the pool indices right-to-left. Only advance to the
1932 next pool when the previous one rolls-over */
1933 for (i=npools-1 ; i >= 0 ; i--) {
1934 pool = PyTuple_GET_ITEM(pools, i);
1935 indices[i]++;
1936 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1937 /* Roll-over and advance to next pool */
1938 indices[i] = 0;
1939 elem = PyTuple_GET_ITEM(pool, 0);
1940 Py_INCREF(elem);
1941 oldelem = PyTuple_GET_ITEM(result, i);
1942 PyTuple_SET_ITEM(result, i, elem);
1943 Py_DECREF(oldelem);
1944 } else {
1945 /* No rollover. Just increment and stop here. */
1946 elem = PyTuple_GET_ITEM(pool, indices[i]);
1947 Py_INCREF(elem);
1948 oldelem = PyTuple_GET_ITEM(result, i);
1949 PyTuple_SET_ITEM(result, i, elem);
1950 Py_DECREF(oldelem);
1951 break;
1952 }
1953 }
1954
1955 /* If i is negative, then the indices have all rolled-over
1956 and we're done. */
1957 if (i < 0)
1958 goto empty;
1959 }
1960
1961 Py_INCREF(result);
1962 return result;
1963
1964 empty:
1965 lz->stopped = 1;
1966 return NULL;
1967 }
1968
1969 PyDoc_STRVAR(product_doc,
1970 "product(*iterables) --> product object\n\
1971 \n\
1972 Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
1973 For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1974 The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1975 cycle in a manner similar to an odometer (with the rightmost element changing\n\
1976 on every iteration).\n\n\
1977 To compute the product of an iterable with itself, specify the number\n\
1978 of repetitions with the optional repeat keyword argument. For example,\n\
1979 product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
1980 product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1981 product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1982
1983 static PyTypeObject product_type = {
1984 PyVarObject_HEAD_INIT(NULL, 0)
1985 "itertools.product", /* tp_name */
1986 sizeof(productobject), /* tp_basicsize */
1987 0, /* tp_itemsize */
1988 /* methods */
1989 (destructor)product_dealloc, /* tp_dealloc */
1990 0, /* tp_print */
1991 0, /* tp_getattr */
1992 0, /* tp_setattr */
1993 0, /* tp_compare */
1994 0, /* tp_repr */
1995 0, /* tp_as_number */
1996 0, /* tp_as_sequence */
1997 0, /* tp_as_mapping */
1998 0, /* tp_hash */
1999 0, /* tp_call */
2000 0, /* tp_str */
2001 PyObject_GenericGetAttr, /* tp_getattro */
2002 0, /* tp_setattro */
2003 0, /* tp_as_buffer */
2004 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2005 Py_TPFLAGS_BASETYPE, /* tp_flags */
2006 product_doc, /* tp_doc */
2007 (traverseproc)product_traverse, /* tp_traverse */
2008 0, /* tp_clear */
2009 0, /* tp_richcompare */
2010 0, /* tp_weaklistoffset */
2011 PyObject_SelfIter, /* tp_iter */
2012 (iternextfunc)product_next, /* tp_iternext */
2013 0, /* tp_methods */
2014 0, /* tp_members */
2015 0, /* tp_getset */
2016 0, /* tp_base */
2017 0, /* tp_dict */
2018 0, /* tp_descr_get */
2019 0, /* tp_descr_set */
2020 0, /* tp_dictoffset */
2021 0, /* tp_init */
2022 0, /* tp_alloc */
2023 product_new, /* tp_new */
2024 PyObject_GC_Del, /* tp_free */
2025 };
2026
2027
2028 /* combinations object ************************************************************/
2029
2030 typedef struct {
2031 PyObject_HEAD
2032 PyObject *pool; /* input converted to a tuple */
2033 Py_ssize_t *indices; /* one index per result element */
2034 PyObject *result; /* most recently returned result tuple */
2035 Py_ssize_t r; /* size of result tuple */
2036 int stopped; /* set to 1 when the combinations iterator is exhausted */
2037 } combinationsobject;
2038
2039 static PyTypeObject combinations_type;
2040
2041 static PyObject *
2042 combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2043 {
2044 combinationsobject *co;
2045 Py_ssize_t n;
2046 Py_ssize_t r;
2047 PyObject *pool = NULL;
2048 PyObject *iterable = NULL;
2049 Py_ssize_t *indices = NULL;
2050 Py_ssize_t i;
2051 static char *kwargs[] = {"iterable", "r", NULL};
2052
2053 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2054 &iterable, &r))
2055 return NULL;
2056
2057 pool = PySequence_Tuple(iterable);
2058 if (pool == NULL)
2059 goto error;
2060 n = PyTuple_GET_SIZE(pool);
2061 if (r < 0) {
2062 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2063 goto error;
2064 }
2065
2066 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2067 if (indices == NULL) {
2068 PyErr_NoMemory();
2069 goto error;
2070 }
2071
2072 for (i=0 ; i<r ; i++)
2073 indices[i] = i;
2074
2075 /* create combinationsobject structure */
2076 co = (combinationsobject *)type->tp_alloc(type, 0);
2077 if (co == NULL)
2078 goto error;
2079
2080 co->pool = pool;
2081 co->indices = indices;
2082 co->result = NULL;
2083 co->r = r;
2084 co->stopped = r > n ? 1 : 0;
2085
2086 return (PyObject *)co;
2087
2088 error:
2089 if (indices != NULL)
2090 PyMem_Free(indices);
2091 Py_XDECREF(pool);
2092 return NULL;
2093 }
2094
2095 static void
2096 combinations_dealloc(combinationsobject *co)
2097 {
2098 PyObject_GC_UnTrack(co);
2099 Py_XDECREF(co->pool);
2100 Py_XDECREF(co->result);
2101 if (co->indices != NULL)
2102 PyMem_Free(co->indices);
2103 Py_TYPE(co)->tp_free(co);
2104 }
2105
2106 static int
2107 combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2108 {
2109 Py_VISIT(co->pool);
2110 Py_VISIT(co->result);
2111 return 0;
2112 }
2113
2114 static PyObject *
2115 combinations_next(combinationsobject *co)
2116 {
2117 PyObject *elem;
2118 PyObject *oldelem;
2119 PyObject *pool = co->pool;
2120 Py_ssize_t *indices = co->indices;
2121 PyObject *result = co->result;
2122 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2123 Py_ssize_t r = co->r;
2124 Py_ssize_t i, j, index;
2125
2126 if (co->stopped)
2127 return NULL;
2128
2129 if (result == NULL) {
2130 /* On the first pass, initialize result tuple using the indices */
2131 result = PyTuple_New(r);
2132 if (result == NULL)
2133 goto empty;
2134 co->result = result;
2135 for (i=0; i<r ; i++) {
2136 index = indices[i];
2137 elem = PyTuple_GET_ITEM(pool, index);
2138 Py_INCREF(elem);
2139 PyTuple_SET_ITEM(result, i, elem);
2140 }
2141 } else {
2142 /* Copy the previous result tuple or re-use it if available */
2143 if (Py_REFCNT(result) > 1) {
2144 PyObject *old_result = result;
2145 result = PyTuple_New(r);
2146 if (result == NULL)
2147 goto empty;
2148 co->result = result;
2149 for (i=0; i<r ; i++) {
2150 elem = PyTuple_GET_ITEM(old_result, i);
2151 Py_INCREF(elem);
2152 PyTuple_SET_ITEM(result, i, elem);
2153 }
2154 Py_DECREF(old_result);
2155 }
2156 /* Now, we've got the only copy so we can update it in-place
2157 * CPython's empty tuple is a singleton and cached in
2158 * PyTuple's freelist.
2159 */
2160 assert(r == 0 || Py_REFCNT(result) == 1);
2161
2162 /* Scan indices right-to-left until finding one that is not
2163 at its maximum (i + n - r). */
2164 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2165 ;
2166
2167 /* If i is negative, then the indices are all at
2168 their maximum value and we're done. */
2169 if (i < 0)
2170 goto empty;
2171
2172 /* Increment the current index which we know is not at its
2173 maximum. Then move back to the right setting each index
2174 to its lowest possible value (one higher than the index
2175 to its left -- this maintains the sort order invariant). */
2176 indices[i]++;
2177 for (j=i+1 ; j<r ; j++)
2178 indices[j] = indices[j-1] + 1;
2179
2180 /* Update the result tuple for the new indices
2181 starting with i, the leftmost index that changed */
2182 for ( ; i<r ; i++) {
2183 index = indices[i];
2184 elem = PyTuple_GET_ITEM(pool, index);
2185 Py_INCREF(elem);
2186 oldelem = PyTuple_GET_ITEM(result, i);
2187 PyTuple_SET_ITEM(result, i, elem);
2188 Py_DECREF(oldelem);
2189 }
2190 }
2191
2192 Py_INCREF(result);
2193 return result;
2194
2195 empty:
2196 co->stopped = 1;
2197 return NULL;
2198 }
2199
2200 PyDoc_STRVAR(combinations_doc,
2201 "combinations(iterable, r) --> combinations object\n\
2202 \n\
2203 Return successive r-length combinations of elements in the iterable.\n\n\
2204 combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2205
2206 static PyTypeObject combinations_type = {
2207 PyVarObject_HEAD_INIT(NULL, 0)
2208 "itertools.combinations", /* tp_name */
2209 sizeof(combinationsobject), /* tp_basicsize */
2210 0, /* tp_itemsize */
2211 /* methods */
2212 (destructor)combinations_dealloc, /* tp_dealloc */
2213 0, /* tp_print */
2214 0, /* tp_getattr */
2215 0, /* tp_setattr */
2216 0, /* tp_compare */
2217 0, /* tp_repr */
2218 0, /* tp_as_number */
2219 0, /* tp_as_sequence */
2220 0, /* tp_as_mapping */
2221 0, /* tp_hash */
2222 0, /* tp_call */
2223 0, /* tp_str */
2224 PyObject_GenericGetAttr, /* tp_getattro */
2225 0, /* tp_setattro */
2226 0, /* tp_as_buffer */
2227 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2228 Py_TPFLAGS_BASETYPE, /* tp_flags */
2229 combinations_doc, /* tp_doc */
2230 (traverseproc)combinations_traverse, /* tp_traverse */
2231 0, /* tp_clear */
2232 0, /* tp_richcompare */
2233 0, /* tp_weaklistoffset */
2234 PyObject_SelfIter, /* tp_iter */
2235 (iternextfunc)combinations_next, /* tp_iternext */
2236 0, /* tp_methods */
2237 0, /* tp_members */
2238 0, /* tp_getset */
2239 0, /* tp_base */
2240 0, /* tp_dict */
2241 0, /* tp_descr_get */
2242 0, /* tp_descr_set */
2243 0, /* tp_dictoffset */
2244 0, /* tp_init */
2245 0, /* tp_alloc */
2246 combinations_new, /* tp_new */
2247 PyObject_GC_Del, /* tp_free */
2248 };
2249
2250
2251 /* combinations with replacement object *******************************************/
2252
2253 /* Equivalent to:
2254
2255 def combinations_with_replacement(iterable, r):
2256 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2257 # number items returned: (n+r-1)! / r! / (n-1)!
2258 pool = tuple(iterable)
2259 n = len(pool)
2260 indices = [0] * r
2261 yield tuple(pool[i] for i in indices)
2262 while 1:
2263 for i in reversed(range(r)):
2264 if indices[i] != n - 1:
2265 break
2266 else:
2267 return
2268 indices[i:] = [indices[i] + 1] * (r - i)
2269 yield tuple(pool[i] for i in indices)
2270
2271 def combinations_with_replacement2(iterable, r):
2272 'Alternate version that filters from product()'
2273 pool = tuple(iterable)
2274 n = len(pool)
2275 for indices in product(range(n), repeat=r):
2276 if sorted(indices) == list(indices):
2277 yield tuple(pool[i] for i in indices)
2278 */
2279 typedef struct {
2280 PyObject_HEAD
2281 PyObject *pool; /* input converted to a tuple */
2282 Py_ssize_t *indices; /* one index per result element */
2283 PyObject *result; /* most recently returned result tuple */
2284 Py_ssize_t r; /* size of result tuple */
2285 int stopped; /* set to 1 when the cwr iterator is exhausted */
2286 } cwrobject;
2287
2288 static PyTypeObject cwr_type;
2289
2290 static PyObject *
2291 cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2292 {
2293 cwrobject *co;
2294 Py_ssize_t n;
2295 Py_ssize_t r;
2296 PyObject *pool = NULL;
2297 PyObject *iterable = NULL;
2298 Py_ssize_t *indices = NULL;
2299 Py_ssize_t i;
2300 static char *kwargs[] = {"iterable", "r", NULL};
2301
2302 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2303 &iterable, &r))
2304 return NULL;
2305
2306 pool = PySequence_Tuple(iterable);
2307 if (pool == NULL)
2308 goto error;
2309 n = PyTuple_GET_SIZE(pool);
2310 if (r < 0) {
2311 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2312 goto error;
2313 }
2314
2315 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2316 if (indices == NULL) {
2317 PyErr_NoMemory();
2318 goto error;
2319 }
2320
2321 for (i=0 ; i<r ; i++)
2322 indices[i] = 0;
2323
2324 /* create cwrobject structure */
2325 co = (cwrobject *)type->tp_alloc(type, 0);
2326 if (co == NULL)
2327 goto error;
2328
2329 co->pool = pool;
2330 co->indices = indices;
2331 co->result = NULL;
2332 co->r = r;
2333 co->stopped = !n && r;
2334
2335 return (PyObject *)co;
2336
2337 error:
2338 if (indices != NULL)
2339 PyMem_Free(indices);
2340 Py_XDECREF(pool);
2341 return NULL;
2342 }
2343
2344 static void
2345 cwr_dealloc(cwrobject *co)
2346 {
2347 PyObject_GC_UnTrack(co);
2348 Py_XDECREF(co->pool);
2349 Py_XDECREF(co->result);
2350 if (co->indices != NULL)
2351 PyMem_Free(co->indices);
2352 Py_TYPE(co)->tp_free(co);
2353 }
2354
2355 static int
2356 cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2357 {
2358 Py_VISIT(co->pool);
2359 Py_VISIT(co->result);
2360 return 0;
2361 }
2362
2363 static PyObject *
2364 cwr_next(cwrobject *co)
2365 {
2366 PyObject *elem;
2367 PyObject *oldelem;
2368 PyObject *pool = co->pool;
2369 Py_ssize_t *indices = co->indices;
2370 PyObject *result = co->result;
2371 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2372 Py_ssize_t r = co->r;
2373 Py_ssize_t i, j, index;
2374
2375 if (co->stopped)
2376 return NULL;
2377
2378 if (result == NULL) {
2379 /* On the first pass, initialize result tuple using the indices */
2380 result = PyTuple_New(r);
2381 if (result == NULL)
2382 goto empty;
2383 co->result = result;
2384 for (i=0; i<r ; i++) {
2385 index = indices[i];
2386 elem = PyTuple_GET_ITEM(pool, index);
2387 Py_INCREF(elem);
2388 PyTuple_SET_ITEM(result, i, elem);
2389 }
2390 } else {
2391 /* Copy the previous result tuple or re-use it if available */
2392 if (Py_REFCNT(result) > 1) {
2393 PyObject *old_result = result;
2394 result = PyTuple_New(r);
2395 if (result == NULL)
2396 goto empty;
2397 co->result = result;
2398 for (i=0; i<r ; i++) {
2399 elem = PyTuple_GET_ITEM(old_result, i);
2400 Py_INCREF(elem);
2401 PyTuple_SET_ITEM(result, i, elem);
2402 }
2403 Py_DECREF(old_result);
2404 }
2405 /* Now, we've got the only copy so we can update it in-place CPython's
2406 empty tuple is a singleton and cached in PyTuple's freelist. */
2407 assert(r == 0 || Py_REFCNT(result) == 1);
2408
2409 /* Scan indices right-to-left until finding one that is not
2410 * at its maximum (n-1). */
2411 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2412 ;
2413
2414 /* If i is negative, then the indices are all at
2415 their maximum value and we're done. */
2416 if (i < 0)
2417 goto empty;
2418
2419 /* Increment the current index which we know is not at its
2420 maximum. Then set all to the right to the same value. */
2421 indices[i]++;
2422 for (j=i+1 ; j<r ; j++)
2423 indices[j] = indices[j-1];
2424
2425 /* Update the result tuple for the new indices
2426 starting with i, the leftmost index that changed */
2427 for ( ; i<r ; i++) {
2428 index = indices[i];
2429 elem = PyTuple_GET_ITEM(pool, index);
2430 Py_INCREF(elem);
2431 oldelem = PyTuple_GET_ITEM(result, i);
2432 PyTuple_SET_ITEM(result, i, elem);
2433 Py_DECREF(oldelem);
2434 }
2435 }
2436
2437 Py_INCREF(result);
2438 return result;
2439
2440 empty:
2441 co->stopped = 1;
2442 return NULL;
2443 }
2444
2445 PyDoc_STRVAR(cwr_doc,
2446 "combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
2447 \n\
2448 Return successive r-length combinations of elements in the iterable\n\
2449 allowing individual elements to have successive repeats.\n\
2450 combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2451
2452 static PyTypeObject cwr_type = {
2453 PyVarObject_HEAD_INIT(NULL, 0)
2454 "itertools.combinations_with_replacement", /* tp_name */
2455 sizeof(cwrobject), /* tp_basicsize */
2456 0, /* tp_itemsize */
2457 /* methods */
2458 (destructor)cwr_dealloc, /* tp_dealloc */
2459 0, /* tp_print */
2460 0, /* tp_getattr */
2461 0, /* tp_setattr */
2462 0, /* tp_compare */
2463 0, /* tp_repr */
2464 0, /* tp_as_number */
2465 0, /* tp_as_sequence */
2466 0, /* tp_as_mapping */
2467 0, /* tp_hash */
2468 0, /* tp_call */
2469 0, /* tp_str */
2470 PyObject_GenericGetAttr, /* tp_getattro */
2471 0, /* tp_setattro */
2472 0, /* tp_as_buffer */
2473 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2474 Py_TPFLAGS_BASETYPE, /* tp_flags */
2475 cwr_doc, /* tp_doc */
2476 (traverseproc)cwr_traverse, /* tp_traverse */
2477 0, /* tp_clear */
2478 0, /* tp_richcompare */
2479 0, /* tp_weaklistoffset */
2480 PyObject_SelfIter, /* tp_iter */
2481 (iternextfunc)cwr_next, /* tp_iternext */
2482 0, /* tp_methods */
2483 0, /* tp_members */
2484 0, /* tp_getset */
2485 0, /* tp_base */
2486 0, /* tp_dict */
2487 0, /* tp_descr_get */
2488 0, /* tp_descr_set */
2489 0, /* tp_dictoffset */
2490 0, /* tp_init */
2491 0, /* tp_alloc */
2492 cwr_new, /* tp_new */
2493 PyObject_GC_Del, /* tp_free */
2494 };
2495
2496
2497 /* permutations object ************************************************************
2498
2499 def permutations(iterable, r=None):
2500 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2501 pool = tuple(iterable)
2502 n = len(pool)
2503 r = n if r is None else r
2504 indices = range(n)
2505 cycles = range(n-r+1, n+1)[::-1]
2506 yield tuple(pool[i] for i in indices[:r])
2507 while n:
2508 for i in reversed(range(r)):
2509 cycles[i] -= 1
2510 if cycles[i] == 0:
2511 indices[i:] = indices[i+1:] + indices[i:i+1]
2512 cycles[i] = n - i
2513 else:
2514 j = cycles[i]
2515 indices[i], indices[-j] = indices[-j], indices[i]
2516 yield tuple(pool[i] for i in indices[:r])
2517 break
2518 else:
2519 return
2520 */
2521
2522 typedef struct {
2523 PyObject_HEAD
2524 PyObject *pool; /* input converted to a tuple */
2525 Py_ssize_t *indices; /* one index per element in the pool */
2526 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2527 PyObject *result; /* most recently returned result tuple */
2528 Py_ssize_t r; /* size of result tuple */
2529 int stopped; /* set to 1 when the permutations iterator is exhausted */
2530 } permutationsobject;
2531
2532 static PyTypeObject permutations_type;
2533
2534 static PyObject *
2535 permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2536 {
2537 permutationsobject *po;
2538 Py_ssize_t n;
2539 Py_ssize_t r;
2540 PyObject *robj = Py_None;
2541 PyObject *pool = NULL;
2542 PyObject *iterable = NULL;
2543 Py_ssize_t *indices = NULL;
2544 Py_ssize_t *cycles = NULL;
2545 Py_ssize_t i;
2546 static char *kwargs[] = {"iterable", "r", NULL};
2547
2548 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2549 &iterable, &robj))
2550 return NULL;
2551
2552 pool = PySequence_Tuple(iterable);
2553 if (pool == NULL)
2554 goto error;
2555 n = PyTuple_GET_SIZE(pool);
2556
2557 r = n;
2558 if (robj != Py_None) {
2559 r = PyInt_AsSsize_t(robj);
2560 if (r == -1 && PyErr_Occurred())
2561 goto error;
2562 }
2563 if (r < 0) {
2564 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2565 goto error;
2566 }
2567
2568 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2569 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2570 if (indices == NULL || cycles == NULL) {
2571 PyErr_NoMemory();
2572 goto error;
2573 }
2574
2575 for (i=0 ; i<n ; i++)
2576 indices[i] = i;
2577 for (i=0 ; i<r ; i++)
2578 cycles[i] = n - i;
2579
2580 /* create permutationsobject structure */
2581 po = (permutationsobject *)type->tp_alloc(type, 0);
2582 if (po == NULL)
2583 goto error;
2584
2585 po->pool = pool;
2586 po->indices = indices;
2587 po->cycles = cycles;
2588 po->result = NULL;
2589 po->r = r;
2590 po->stopped = r > n ? 1 : 0;
2591
2592 return (PyObject *)po;
2593
2594 error:
2595 if (indices != NULL)
2596 PyMem_Free(indices);
2597 if (cycles != NULL)
2598 PyMem_Free(cycles);
2599 Py_XDECREF(pool);
2600 return NULL;
2601 }
2602
2603 static void
2604 permutations_dealloc(permutationsobject *po)
2605 {
2606 PyObject_GC_UnTrack(po);
2607 Py_XDECREF(po->pool);
2608 Py_XDECREF(po->result);
2609 PyMem_Free(po->indices);
2610 PyMem_Free(po->cycles);
2611 Py_TYPE(po)->tp_free(po);
2612 }
2613
2614 static int
2615 permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2616 {
2617 Py_VISIT(po->pool);
2618 Py_VISIT(po->result);
2619 return 0;
2620 }
2621
2622 static PyObject *
2623 permutations_next(permutationsobject *po)
2624 {
2625 PyObject *elem;
2626 PyObject *oldelem;
2627 PyObject *pool = po->pool;
2628 Py_ssize_t *indices = po->indices;
2629 Py_ssize_t *cycles = po->cycles;
2630 PyObject *result = po->result;
2631 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2632 Py_ssize_t r = po->r;
2633 Py_ssize_t i, j, k, index;
2634
2635 if (po->stopped)
2636 return NULL;
2637
2638 if (result == NULL) {
2639 /* On the first pass, initialize result tuple using the indices */
2640 result = PyTuple_New(r);
2641 if (result == NULL)
2642 goto empty;
2643 po->result = result;
2644 for (i=0; i<r ; i++) {
2645 index = indices[i];
2646 elem = PyTuple_GET_ITEM(pool, index);
2647 Py_INCREF(elem);
2648 PyTuple_SET_ITEM(result, i, elem);
2649 }
2650 } else {
2651 if (n == 0)
2652 goto empty;
2653
2654 /* Copy the previous result tuple or re-use it if available */
2655 if (Py_REFCNT(result) > 1) {
2656 PyObject *old_result = result;
2657 result = PyTuple_New(r);
2658 if (result == NULL)
2659 goto empty;
2660 po->result = result;
2661 for (i=0; i<r ; i++) {
2662 elem = PyTuple_GET_ITEM(old_result, i);
2663 Py_INCREF(elem);
2664 PyTuple_SET_ITEM(result, i, elem);
2665 }
2666 Py_DECREF(old_result);
2667 }
2668 /* Now, we've got the only copy so we can update it in-place */
2669 assert(r == 0 || Py_REFCNT(result) == 1);
2670
2671 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2672 for (i=r-1 ; i>=0 ; i--) {
2673 cycles[i] -= 1;
2674 if (cycles[i] == 0) {
2675 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2676 index = indices[i];
2677 for (j=i ; j<n-1 ; j++)
2678 indices[j] = indices[j+1];
2679 indices[n-1] = index;
2680 cycles[i] = n - i;
2681 } else {
2682 j = cycles[i];
2683 index = indices[i];
2684 indices[i] = indices[n-j];
2685 indices[n-j] = index;
2686
2687 for (k=i; k<r ; k++) {
2688 /* start with i, the leftmost element that changed */
2689 /* yield tuple(pool[k] for k in indices[:r]) */
2690 index = indices[k];
2691 elem = PyTuple_GET_ITEM(pool, index);
2692 Py_INCREF(elem);
2693 oldelem = PyTuple_GET_ITEM(result, k);
2694 PyTuple_SET_ITEM(result, k, elem);
2695 Py_DECREF(oldelem);
2696 }
2697 break;
2698 }
2699 }
2700 /* If i is negative, then the cycles have all
2701 rolled-over and we're done. */
2702 if (i < 0)
2703 goto empty;
2704 }
2705 Py_INCREF(result);
2706 return result;
2707
2708 empty:
2709 po->stopped = 1;
2710 return NULL;
2711 }
2712
2713 PyDoc_STRVAR(permutations_doc,
2714 "permutations(iterable[, r]) --> permutations object\n\
2715 \n\
2716 Return successive r-length permutations of elements in the iterable.\n\n\
2717 permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
2718
2719 static PyTypeObject permutations_type = {
2720 PyVarObject_HEAD_INIT(NULL, 0)
2721 "itertools.permutations", /* tp_name */
2722 sizeof(permutationsobject), /* tp_basicsize */
2723 0, /* tp_itemsize */
2724 /* methods */
2725 (destructor)permutations_dealloc, /* tp_dealloc */
2726 0, /* tp_print */
2727 0, /* tp_getattr */
2728 0, /* tp_setattr */
2729 0, /* tp_compare */
2730 0, /* tp_repr */
2731 0, /* tp_as_number */
2732 0, /* tp_as_sequence */
2733 0, /* tp_as_mapping */
2734 0, /* tp_hash */
2735 0, /* tp_call */
2736 0, /* tp_str */
2737 PyObject_GenericGetAttr, /* tp_getattro */
2738 0, /* tp_setattro */
2739 0, /* tp_as_buffer */
2740 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2741 Py_TPFLAGS_BASETYPE, /* tp_flags */
2742 permutations_doc, /* tp_doc */
2743 (traverseproc)permutations_traverse, /* tp_traverse */
2744 0, /* tp_clear */
2745 0, /* tp_richcompare */
2746 0, /* tp_weaklistoffset */
2747 PyObject_SelfIter, /* tp_iter */
2748 (iternextfunc)permutations_next, /* tp_iternext */
2749 0, /* tp_methods */
2750 0, /* tp_members */
2751 0, /* tp_getset */
2752 0, /* tp_base */
2753 0, /* tp_dict */
2754 0, /* tp_descr_get */
2755 0, /* tp_descr_set */
2756 0, /* tp_dictoffset */
2757 0, /* tp_init */
2758 0, /* tp_alloc */
2759 permutations_new, /* tp_new */
2760 PyObject_GC_Del, /* tp_free */
2761 };
2762
2763
2764 /* compress object ************************************************************/
2765
2766 /* Equivalent to:
2767
2768 def compress(data, selectors):
2769 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2770 return (d for d, s in izip(data, selectors) if s)
2771 */
2772
2773 typedef struct {
2774 PyObject_HEAD
2775 PyObject *data;
2776 PyObject *selectors;
2777 } compressobject;
2778
2779 static PyTypeObject compress_type;
2780
2781 static PyObject *
2782 compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2783 {
2784 PyObject *seq1, *seq2;
2785 PyObject *data=NULL, *selectors=NULL;
2786 compressobject *lz;
2787 static char *kwargs[] = {"data", "selectors", NULL};
2788
2789 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2790 return NULL;
2791
2792 data = PyObject_GetIter(seq1);
2793 if (data == NULL)
2794 goto fail;
2795 selectors = PyObject_GetIter(seq2);
2796 if (selectors == NULL)
2797 goto fail;
2798
2799 /* create compressobject structure */
2800 lz = (compressobject *)type->tp_alloc(type, 0);
2801 if (lz == NULL)
2802 goto fail;
2803 lz->data = data;
2804 lz->selectors = selectors;
2805 return (PyObject *)lz;
2806
2807 fail:
2808 Py_XDECREF(data);
2809 Py_XDECREF(selectors);
2810 return NULL;
2811 }
2812
2813 static void
2814 compress_dealloc(compressobject *lz)
2815 {
2816 PyObject_GC_UnTrack(lz);
2817 Py_XDECREF(lz->data);
2818 Py_XDECREF(lz->selectors);
2819 Py_TYPE(lz)->tp_free(lz);
2820 }
2821
2822 static int
2823 compress_traverse(compressobject *lz, visitproc visit, void *arg)
2824 {
2825 Py_VISIT(lz->data);
2826 Py_VISIT(lz->selectors);
2827 return 0;
2828 }
2829
2830 static PyObject *
2831 compress_next(compressobject *lz)
2832 {
2833 PyObject *data = lz->data, *selectors = lz->selectors;
2834 PyObject *datum, *selector;
2835 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2836 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2837 int ok;
2838
2839 while (1) {
2840 /* Steps: get datum, get selector, evaluate selector.
2841 Order is important (to match the pure python version
2842 in terms of which input gets a chance to raise an
2843 exception first).
2844 */
2845
2846 datum = datanext(data);
2847 if (datum == NULL)
2848 return NULL;
2849
2850 selector = selectornext(selectors);
2851 if (selector == NULL) {
2852 Py_DECREF(datum);
2853 return NULL;
2854 }
2855
2856 ok = PyObject_IsTrue(selector);
2857 Py_DECREF(selector);
2858 if (ok == 1)
2859 return datum;
2860 Py_DECREF(datum);
2861 if (ok == -1)
2862 return NULL;
2863 }
2864 }
2865
2866 PyDoc_STRVAR(compress_doc,
2867 "compress(data, selectors) --> iterator over selected data\n\
2868 \n\
2869 Return data elements corresponding to true selector elements.\n\
2870 Forms a shorter iterator from selected data elements using the\n\
2871 selectors to choose the data elements.");
2872
2873 static PyTypeObject compress_type = {
2874 PyVarObject_HEAD_INIT(NULL, 0)
2875 "itertools.compress", /* tp_name */
2876 sizeof(compressobject), /* tp_basicsize */
2877 0, /* tp_itemsize */
2878 /* methods */
2879 (destructor)compress_dealloc, /* tp_dealloc */
2880 0, /* tp_print */
2881 0, /* tp_getattr */
2882 0, /* tp_setattr */
2883 0, /* tp_compare */
2884 0, /* tp_repr */
2885 0, /* tp_as_number */
2886 0, /* tp_as_sequence */
2887 0, /* tp_as_mapping */
2888 0, /* tp_hash */
2889 0, /* tp_call */
2890 0, /* tp_str */
2891 PyObject_GenericGetAttr, /* tp_getattro */
2892 0, /* tp_setattro */
2893 0, /* tp_as_buffer */
2894 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2895 Py_TPFLAGS_BASETYPE, /* tp_flags */
2896 compress_doc, /* tp_doc */
2897 (traverseproc)compress_traverse, /* tp_traverse */
2898 0, /* tp_clear */
2899 0, /* tp_richcompare */
2900 0, /* tp_weaklistoffset */
2901 PyObject_SelfIter, /* tp_iter */
2902 (iternextfunc)compress_next, /* tp_iternext */
2903 0, /* tp_methods */
2904 0, /* tp_members */
2905 0, /* tp_getset */
2906 0, /* tp_base */
2907 0, /* tp_dict */
2908 0, /* tp_descr_get */
2909 0, /* tp_descr_set */
2910 0, /* tp_dictoffset */
2911 0, /* tp_init */
2912 0, /* tp_alloc */
2913 compress_new, /* tp_new */
2914 PyObject_GC_Del, /* tp_free */
2915 };
2916
2917
2918 /* ifilter object ************************************************************/
2919
2920 typedef struct {
2921 PyObject_HEAD
2922 PyObject *func;
2923 PyObject *it;
2924 } ifilterobject;
2925
2926 static PyTypeObject ifilter_type;
2927
2928 static PyObject *
2929 ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2930 {
2931 PyObject *func, *seq;
2932 PyObject *it;
2933 ifilterobject *lz;
2934
2935 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2936 return NULL;
2937
2938 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2939 return NULL;
2940
2941 /* Get iterator. */
2942 it = PyObject_GetIter(seq);
2943 if (it == NULL)
2944 return NULL;
2945
2946 /* create ifilterobject structure */
2947 lz = (ifilterobject *)type->tp_alloc(type, 0);
2948 if (lz == NULL) {
2949 Py_DECREF(it);
2950 return NULL;
2951 }
2952 Py_INCREF(func);
2953 lz->func = func;
2954 lz->it = it;
2955
2956 return (PyObject *)lz;
2957 }
2958
2959 static void
2960 ifilter_dealloc(ifilterobject *lz)
2961 {
2962 PyObject_GC_UnTrack(lz);
2963 Py_XDECREF(lz->func);
2964 Py_XDECREF(lz->it);
2965 Py_TYPE(lz)->tp_free(lz);
2966 }
2967
2968 static int
2969 ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2970 {
2971 Py_VISIT(lz->it);
2972 Py_VISIT(lz->func);
2973 return 0;
2974 }
2975
2976 static PyObject *
2977 ifilter_next(ifilterobject *lz)
2978 {
2979 PyObject *item;
2980 PyObject *it = lz->it;
2981 long ok;
2982 PyObject *(*iternext)(PyObject *);
2983
2984 iternext = *Py_TYPE(it)->tp_iternext;
2985 for (;;) {
2986 item = iternext(it);
2987 if (item == NULL)
2988 return NULL;
2989
2990 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2991 ok = PyObject_IsTrue(item);
2992 } else {
2993 PyObject *good;
2994 good = PyObject_CallFunctionObjArgs(lz->func,
2995 item, NULL);
2996 if (good == NULL) {
2997 Py_DECREF(item);
2998 return NULL;
2999 }
3000 ok = PyObject_IsTrue(good);
3001 Py_DECREF(good);
3002 }
3003 if (ok)
3004 return item;
3005 Py_DECREF(item);
3006 }
3007 }
3008
3009 PyDoc_STRVAR(ifilter_doc,
3010 "ifilter(function or None, sequence) --> ifilter object\n\
3011 \n\
3012 Return those items of sequence for which function(item) is true.\n\
3013 If function is None, return the items that are true.");
3014
3015 static PyTypeObject ifilter_type = {
3016 PyVarObject_HEAD_INIT(NULL, 0)
3017 "itertools.ifilter", /* tp_name */
3018 sizeof(ifilterobject), /* tp_basicsize */
3019 0, /* tp_itemsize */
3020 /* methods */
3021 (destructor)ifilter_dealloc, /* tp_dealloc */
3022 0, /* tp_print */
3023 0, /* tp_getattr */
3024 0, /* tp_setattr */
3025 0, /* tp_compare */
3026 0, /* tp_repr */
3027 0, /* tp_as_number */
3028 0, /* tp_as_sequence */
3029 0, /* tp_as_mapping */
3030 0, /* tp_hash */
3031 0, /* tp_call */
3032 0, /* tp_str */
3033 PyObject_GenericGetAttr, /* tp_getattro */
3034 0, /* tp_setattro */
3035 0, /* tp_as_buffer */
3036 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3037 Py_TPFLAGS_BASETYPE, /* tp_flags */
3038 ifilter_doc, /* tp_doc */
3039 (traverseproc)ifilter_traverse, /* tp_traverse */
3040 0, /* tp_clear */
3041 0, /* tp_richcompare */
3042 0, /* tp_weaklistoffset */
3043 PyObject_SelfIter, /* tp_iter */
3044 (iternextfunc)ifilter_next, /* tp_iternext */
3045 0, /* tp_methods */
3046 0, /* tp_members */
3047 0, /* tp_getset */
3048 0, /* tp_base */
3049 0, /* tp_dict */
3050 0, /* tp_descr_get */
3051 0, /* tp_descr_set */
3052 0, /* tp_dictoffset */
3053 0, /* tp_init */
3054 0, /* tp_alloc */
3055 ifilter_new, /* tp_new */
3056 PyObject_GC_Del, /* tp_free */
3057 };
3058
3059
3060 /* ifilterfalse object ************************************************************/
3061
3062 typedef struct {
3063 PyObject_HEAD
3064 PyObject *func;
3065 PyObject *it;
3066 } ifilterfalseobject;
3067
3068 static PyTypeObject ifilterfalse_type;
3069
3070 static PyObject *
3071 ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3072 {
3073 PyObject *func, *seq;
3074 PyObject *it;
3075 ifilterfalseobject *lz;
3076
3077 if (type == &ifilterfalse_type &&
3078 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3079 return NULL;
3080
3081 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3082 return NULL;
3083
3084 /* Get iterator. */
3085 it = PyObject_GetIter(seq);
3086 if (it == NULL)
3087 return NULL;
3088
3089 /* create ifilterfalseobject structure */
3090 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3091 if (lz == NULL) {
3092 Py_DECREF(it);
3093 return NULL;
3094 }
3095 Py_INCREF(func);
3096 lz->func = func;
3097 lz->it = it;
3098
3099 return (PyObject *)lz;
3100 }
3101
3102 static void
3103 ifilterfalse_dealloc(ifilterfalseobject *lz)
3104 {
3105 PyObject_GC_UnTrack(lz);
3106 Py_XDECREF(lz->func);
3107 Py_XDECREF(lz->it);
3108 Py_TYPE(lz)->tp_free(lz);
3109 }
3110
3111 static int
3112 ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3113 {
3114 Py_VISIT(lz->it);
3115 Py_VISIT(lz->func);
3116 return 0;
3117 }
3118
3119 static PyObject *
3120 ifilterfalse_next(ifilterfalseobject *lz)
3121 {
3122 PyObject *item;
3123 PyObject *it = lz->it;
3124 long ok;
3125 PyObject *(*iternext)(PyObject *);
3126
3127 iternext = *Py_TYPE(it)->tp_iternext;
3128 for (;;) {
3129 item = iternext(it);
3130 if (item == NULL)
3131 return NULL;
3132
3133 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3134 ok = PyObject_IsTrue(item);
3135 } else {
3136 PyObject *good;
3137 good = PyObject_CallFunctionObjArgs(lz->func,
3138 item, NULL);
3139 if (good == NULL) {
3140 Py_DECREF(item);
3141 return NULL;
3142 }
3143 ok = PyObject_IsTrue(good);
3144 Py_DECREF(good);
3145 }
3146 if (!ok)
3147 return item;
3148 Py_DECREF(item);
3149 }
3150 }
3151
3152 PyDoc_STRVAR(ifilterfalse_doc,
3153 "ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3154 \n\
3155 Return those items of sequence for which function(item) is false.\n\
3156 If function is None, return the items that are false.");
3157
3158 static PyTypeObject ifilterfalse_type = {
3159 PyVarObject_HEAD_INIT(NULL, 0)
3160 "itertools.ifilterfalse", /* tp_name */
3161 sizeof(ifilterfalseobject), /* tp_basicsize */
3162 0, /* tp_itemsize */
3163 /* methods */
3164 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3165 0, /* tp_print */
3166 0, /* tp_getattr */
3167 0, /* tp_setattr */
3168 0, /* tp_compare */
3169 0, /* tp_repr */
3170 0, /* tp_as_number */
3171 0, /* tp_as_sequence */
3172 0, /* tp_as_mapping */
3173 0, /* tp_hash */
3174 0, /* tp_call */
3175 0, /* tp_str */
3176 PyObject_GenericGetAttr, /* tp_getattro */
3177 0, /* tp_setattro */
3178 0, /* tp_as_buffer */
3179 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3180 Py_TPFLAGS_BASETYPE, /* tp_flags */
3181 ifilterfalse_doc, /* tp_doc */
3182 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3183 0, /* tp_clear */
3184 0, /* tp_richcompare */
3185 0, /* tp_weaklistoffset */
3186 PyObject_SelfIter, /* tp_iter */
3187 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3188 0, /* tp_methods */
3189 0, /* tp_members */
3190 0, /* tp_getset */
3191 0, /* tp_base */
3192 0, /* tp_dict */
3193 0, /* tp_descr_get */
3194 0, /* tp_descr_set */
3195 0, /* tp_dictoffset */
3196 0, /* tp_init */
3197 0, /* tp_alloc */
3198 ifilterfalse_new, /* tp_new */
3199 PyObject_GC_Del, /* tp_free */
3200 };
3201
3202
3203 /* count object ************************************************************/
3204
3205 typedef struct {
3206 PyObject_HEAD
3207 Py_ssize_t cnt;
3208 PyObject *long_cnt;
3209 PyObject *long_step;
3210 } countobject;
3211
3212 /* Counting logic and invariants:
3213
3214 fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3215
3216 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3217 Advances with: cnt += 1
3218 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
3219
3220 slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3221
3222 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3223 All counting is done with python objects (no overflows or underflows).
3224 Advances with: long_cnt += long_step
3225 Step may be zero -- effectively a slow version of repeat(cnt).
3226 Either long_cnt or long_step may be a float, Fraction, or Decimal.
3227 */
3228
3229 static PyTypeObject count_type;
3230
3231 static PyObject *
3232 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3233 {
3234 countobject *lz;
3235 int slow_mode = 0;
3236 Py_ssize_t cnt = 0;
3237 PyObject *long_cnt = NULL;
3238 PyObject *long_step = NULL;
3239 static char *kwlist[] = {"start", "step", 0};
3240
3241 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3242 kwlist, &long_cnt, &long_step))
3243 return NULL;
3244
3245 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3246 (long_step != NULL && !PyNumber_Check(long_step))) {
3247 PyErr_SetString(PyExc_TypeError, "a number is required");
3248 return NULL;
3249 }
3250
3251 if (long_cnt != NULL) {
3252 cnt = PyInt_AsSsize_t(long_cnt);
3253 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3254 PyErr_Clear();
3255 slow_mode = 1;
3256 }
3257 Py_INCREF(long_cnt);
3258 } else {
3259 cnt = 0;
3260 long_cnt = PyInt_FromLong(0);
3261 }
3262
3263 /* If not specified, step defaults to 1 */
3264 if (long_step == NULL) {
3265 long_step = PyInt_FromLong(1);
3266 if (long_step == NULL) {
3267 Py_DECREF(long_cnt);
3268 return NULL;
3269 }
3270 } else
3271 Py_INCREF(long_step);
3272
3273 assert(long_cnt != NULL && long_step != NULL);
3274
3275 /* Fast mode only works when the step is 1 */
3276 if (!PyInt_Check(long_step) ||
3277 PyInt_AS_LONG(long_step) != 1) {
3278 slow_mode = 1;
3279 }
3280
3281 if (slow_mode)
3282 cnt = PY_SSIZE_T_MAX;
3283 else
3284 Py_CLEAR(long_cnt);
3285
3286 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3287 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3288 assert(slow_mode ||
3289 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
3290
3291 /* create countobject structure */
3292 lz = (countobject *)type->tp_alloc(type, 0);
3293 if (lz == NULL) {
3294 Py_XDECREF(long_cnt);
3295 return NULL;
3296 }
3297 lz->cnt = cnt;
3298 lz->long_cnt = long_cnt;
3299 lz->long_step = long_step;
3300
3301 return (PyObject *)lz;
3302 }
3303
3304 static void
3305 count_dealloc(countobject *lz)
3306 {
3307 PyObject_GC_UnTrack(lz);
3308 Py_XDECREF(lz->long_cnt);
3309 Py_XDECREF(lz->long_step);
3310 Py_TYPE(lz)->tp_free(lz);
3311 }
3312
3313 static int
3314 count_traverse(countobject *lz, visitproc visit, void *arg)
3315 {
3316 Py_VISIT(lz->long_cnt);
3317 Py_VISIT(lz->long_step);
3318 return 0;
3319 }
3320
3321 static PyObject *
3322 count_nextlong(countobject *lz)
3323 {
3324 PyObject *long_cnt;
3325 PyObject *stepped_up;
3326
3327 long_cnt = lz->long_cnt;
3328 if (long_cnt == NULL) {
3329 /* Switch to slow_mode */
3330 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3331 if (long_cnt == NULL)
3332 return NULL;
3333 }
3334 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
3335
3336 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3337 if (stepped_up == NULL)
3338 return NULL;
3339 lz->long_cnt = stepped_up;
3340 return long_cnt;
3341 }
3342
3343 static PyObject *
3344 count_next(countobject *lz)
3345 {
3346 if (lz->cnt == PY_SSIZE_T_MAX)
3347 return count_nextlong(lz);
3348 return PyInt_FromSsize_t(lz->cnt++);
3349 }
3350
3351 static PyObject *
3352 count_repr(countobject *lz)
3353 {
3354 PyObject *cnt_repr, *step_repr = NULL;
3355 PyObject *result = NULL;
3356
3357 if (lz->cnt != PY_SSIZE_T_MAX)
3358 return PyString_FromFormat("count(%zd)", lz->cnt);
3359
3360 cnt_repr = PyObject_Repr(lz->long_cnt);
3361 if (cnt_repr == NULL)
3362 return NULL;
3363
3364 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3365 /* Don't display step when it is an integer equal to 1 */
3366 result = PyString_FromFormat("count(%s)",
3367 PyString_AS_STRING(cnt_repr));
3368 } else {
3369 step_repr = PyObject_Repr(lz->long_step);
3370 if (step_repr != NULL)
3371 result = PyString_FromFormat("count(%s, %s)",
3372 PyString_AS_STRING(cnt_repr),
3373 PyString_AS_STRING(step_repr));
3374 }
3375 Py_DECREF(cnt_repr);
3376 Py_XDECREF(step_repr);
3377 return result;
3378 }
3379
3380 static PyObject *
3381 count_reduce(countobject *lz)
3382 {
3383 if (lz->cnt == PY_SSIZE_T_MAX)
3384 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3385 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
3386 }
3387
3388 PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3389
3390 static PyMethodDef count_methods[] = {
3391 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3392 count_reduce_doc},
3393 {NULL, NULL} /* sentinel */
3394 };
3395
3396 PyDoc_STRVAR(count_doc,
3397 "count(start=0, step=1) --> count object\n\
3398 \n\
3399 Return a count object whose .next() method returns consecutive values.\n\
3400 Equivalent to:\n\n\
3401 def count(firstval=0, step=1):\n\
3402 x = firstval\n\
3403 while 1:\n\
3404 yield x\n\
3405 x += step\n");
3406
3407 static PyTypeObject count_type = {
3408 PyVarObject_HEAD_INIT(NULL, 0)
3409 "itertools.count", /* tp_name */
3410 sizeof(countobject), /* tp_basicsize */
3411 0, /* tp_itemsize */
3412 /* methods */
3413 (destructor)count_dealloc, /* tp_dealloc */
3414 0, /* tp_print */
3415 0, /* tp_getattr */
3416 0, /* tp_setattr */
3417 0, /* tp_compare */
3418 (reprfunc)count_repr, /* tp_repr */
3419 0, /* tp_as_number */
3420 0, /* tp_as_sequence */
3421 0, /* tp_as_mapping */
3422 0, /* tp_hash */
3423 0, /* tp_call */
3424 0, /* tp_str */
3425 PyObject_GenericGetAttr, /* tp_getattro */
3426 0, /* tp_setattro */
3427 0, /* tp_as_buffer */
3428 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3429 Py_TPFLAGS_BASETYPE, /* tp_flags */
3430 count_doc, /* tp_doc */
3431 (traverseproc)count_traverse, /* tp_traverse */
3432 0, /* tp_clear */
3433 0, /* tp_richcompare */
3434 0, /* tp_weaklistoffset */
3435 PyObject_SelfIter, /* tp_iter */
3436 (iternextfunc)count_next, /* tp_iternext */
3437 count_methods, /* tp_methods */
3438 0, /* tp_members */
3439 0, /* tp_getset */
3440 0, /* tp_base */
3441 0, /* tp_dict */
3442 0, /* tp_descr_get */
3443 0, /* tp_descr_set */
3444 0, /* tp_dictoffset */
3445 0, /* tp_init */
3446 0, /* tp_alloc */
3447 count_new, /* tp_new */
3448 PyObject_GC_Del, /* tp_free */
3449 };
3450
3451
3452 /* izip object ************************************************************/
3453
3454 #include "Python.h"
3455
3456 typedef struct {
3457 PyObject_HEAD
3458 Py_ssize_t tuplesize;
3459 PyObject *ittuple; /* tuple of iterators */
3460 PyObject *result;
3461 } izipobject;
3462
3463 static PyTypeObject izip_type;
3464
3465 static PyObject *
3466 izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3467 {
3468 izipobject *lz;
3469 Py_ssize_t i;
3470 PyObject *ittuple; /* tuple of iterators */
3471 PyObject *result;
3472 Py_ssize_t tuplesize = PySequence_Length(args);
3473
3474 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3475 return NULL;
3476
3477 /* args must be a tuple */
3478 assert(PyTuple_Check(args));
3479
3480 /* obtain iterators */
3481 ittuple = PyTuple_New(tuplesize);
3482 if (ittuple == NULL)
3483 return NULL;
3484 for (i=0; i < tuplesize; ++i) {
3485 PyObject *item = PyTuple_GET_ITEM(args, i);
3486 PyObject *it = PyObject_GetIter(item);
3487 if (it == NULL) {
3488 if (PyErr_ExceptionMatches(PyExc_TypeError))
3489 PyErr_Format(PyExc_TypeError,
3490 "izip argument #%zd must support iteration",
3491 i+1);
3492 Py_DECREF(ittuple);
3493 return NULL;
3494 }
3495 PyTuple_SET_ITEM(ittuple, i, it);
3496 }
3497
3498 /* create a result holder */
3499 result = PyTuple_New(tuplesize);
3500 if (result == NULL) {
3501 Py_DECREF(ittuple);
3502 return NULL;
3503 }
3504 for (i=0 ; i < tuplesize ; i++) {
3505 Py_INCREF(Py_None);
3506 PyTuple_SET_ITEM(result, i, Py_None);
3507 }
3508
3509 /* create izipobject structure */
3510 lz = (izipobject *)type->tp_alloc(type, 0);
3511 if (lz == NULL) {
3512 Py_DECREF(ittuple);
3513 Py_DECREF(result);
3514 return NULL;
3515 }
3516 lz->ittuple = ittuple;
3517 lz->tuplesize = tuplesize;
3518 lz->result = result;
3519
3520 return (PyObject *)lz;
3521 }
3522
3523 static void
3524 izip_dealloc(izipobject *lz)
3525 {
3526 PyObject_GC_UnTrack(lz);
3527 Py_XDECREF(lz->ittuple);
3528 Py_XDECREF(lz->result);
3529 Py_TYPE(lz)->tp_free(lz);
3530 }
3531
3532 static int
3533 izip_traverse(izipobject *lz, visitproc visit, void *arg)
3534 {
3535 Py_VISIT(lz->ittuple);
3536 Py_VISIT(lz->result);
3537 return 0;
3538 }
3539
3540 static PyObject *
3541 izip_next(izipobject *lz)
3542 {
3543 Py_ssize_t i;
3544 Py_ssize_t tuplesize = lz->tuplesize;
3545 PyObject *result = lz->result;
3546 PyObject *it;
3547 PyObject *item;
3548 PyObject *olditem;
3549
3550 if (tuplesize == 0)
3551 return NULL;
3552 if (Py_REFCNT(result) == 1) {
3553 Py_INCREF(result);
3554 for (i=0 ; i < tuplesize ; i++) {
3555 it = PyTuple_GET_ITEM(lz->ittuple, i);
3556 item = (*Py_TYPE(it)->tp_iternext)(it);
3557 if (item == NULL) {
3558 Py_DECREF(result);
3559 return NULL;
3560 }
3561 olditem = PyTuple_GET_ITEM(result, i);
3562 PyTuple_SET_ITEM(result, i, item);
3563 Py_DECREF(olditem);
3564 }
3565 } else {
3566 result = PyTuple_New(tuplesize);
3567 if (result == NULL)
3568 return NULL;
3569 for (i=0 ; i < tuplesize ; i++) {
3570 it = PyTuple_GET_ITEM(lz->ittuple, i);
3571 item = (*Py_TYPE(it)->tp_iternext)(it);
3572 if (item == NULL) {
3573 Py_DECREF(result);
3574 return NULL;
3575 }
3576 PyTuple_SET_ITEM(result, i, item);
3577 }
3578 }
3579 return result;
3580 }
3581
3582 PyDoc_STRVAR(izip_doc,
3583 "izip(iter1 [,iter2 [...]]) --> izip object\n\
3584 \n\
3585 Return a izip object whose .next() method returns a tuple where\n\
3586 the i-th element comes from the i-th iterable argument. The .next()\n\
3587 method continues until the shortest iterable in the argument sequence\n\
3588 is exhausted and then it raises StopIteration. Works like the zip()\n\
3589 function but consumes less memory by returning an iterator instead of\n\
3590 a list.");
3591
3592 static PyTypeObject izip_type = {
3593 PyVarObject_HEAD_INIT(NULL, 0)
3594 "itertools.izip", /* tp_name */
3595 sizeof(izipobject), /* tp_basicsize */
3596 0, /* tp_itemsize */
3597 /* methods */
3598 (destructor)izip_dealloc, /* tp_dealloc */
3599 0, /* tp_print */
3600 0, /* tp_getattr */
3601 0, /* tp_setattr */
3602 0, /* tp_compare */
3603 0, /* tp_repr */
3604 0, /* tp_as_number */
3605 0, /* tp_as_sequence */
3606 0, /* tp_as_mapping */
3607 0, /* tp_hash */
3608 0, /* tp_call */
3609 0, /* tp_str */
3610 PyObject_GenericGetAttr, /* tp_getattro */
3611 0, /* tp_setattro */
3612 0, /* tp_as_buffer */
3613 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3614 Py_TPFLAGS_BASETYPE, /* tp_flags */
3615 izip_doc, /* tp_doc */
3616 (traverseproc)izip_traverse, /* tp_traverse */
3617 0, /* tp_clear */
3618 0, /* tp_richcompare */
3619 0, /* tp_weaklistoffset */
3620 PyObject_SelfIter, /* tp_iter */
3621 (iternextfunc)izip_next, /* tp_iternext */
3622 0, /* tp_methods */
3623 0, /* tp_members */
3624 0, /* tp_getset */
3625 0, /* tp_base */
3626 0, /* tp_dict */
3627 0, /* tp_descr_get */
3628 0, /* tp_descr_set */
3629 0, /* tp_dictoffset */
3630 0, /* tp_init */
3631 0, /* tp_alloc */
3632 izip_new, /* tp_new */
3633 PyObject_GC_Del, /* tp_free */
3634 };
3635
3636
3637 /* repeat object ************************************************************/
3638
3639 typedef struct {
3640 PyObject_HEAD
3641 PyObject *element;
3642 Py_ssize_t cnt;
3643 } repeatobject;
3644
3645 static PyTypeObject repeat_type;
3646
3647 static PyObject *
3648 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3649 {
3650 repeatobject *ro;
3651 PyObject *element;
3652 Py_ssize_t cnt = -1;
3653 static char *kwargs[] = {"object", "times", NULL};
3654
3655 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3656 &element, &cnt))
3657 return NULL;
3658
3659 if (PyTuple_Size(args) == 2 && cnt < 0)
3660 cnt = 0;
3661
3662 ro = (repeatobject *)type->tp_alloc(type, 0);
3663 if (ro == NULL)
3664 return NULL;
3665 Py_INCREF(element);
3666 ro->element = element;
3667 ro->cnt = cnt;
3668 return (PyObject *)ro;
3669 }
3670
3671 static void
3672 repeat_dealloc(repeatobject *ro)
3673 {
3674 PyObject_GC_UnTrack(ro);
3675 Py_XDECREF(ro->element);
3676 Py_TYPE(ro)->tp_free(ro);
3677 }
3678
3679 static int
3680 repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3681 {
3682 Py_VISIT(ro->element);
3683 return 0;
3684 }
3685
3686 static PyObject *
3687 repeat_next(repeatobject *ro)
3688 {
3689 if (ro->cnt == 0)
3690 return NULL;
3691 if (ro->cnt > 0)
3692 ro->cnt--;
3693 Py_INCREF(ro->element);
3694 return ro->element;
3695 }
3696
3697 static PyObject *
3698 repeat_repr(repeatobject *ro)
3699 {
3700 PyObject *result, *objrepr;
3701
3702 objrepr = PyObject_Repr(ro->element);
3703 if (objrepr == NULL)
3704 return NULL;
3705
3706 if (ro->cnt == -1)
3707 result = PyString_FromFormat("repeat(%s)",
3708 PyString_AS_STRING(objrepr));
3709 else
3710 result = PyString_FromFormat("repeat(%s, %zd)",
3711 PyString_AS_STRING(objrepr), ro->cnt);
3712 Py_DECREF(objrepr);
3713 return result;
3714 }
3715
3716 static PyObject *
3717 repeat_len(repeatobject *ro)
3718 {
3719 if (ro->cnt == -1) {
3720 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3721 return NULL;
3722 }
3723 return PyInt_FromSize_t(ro->cnt);
3724 }
3725
3726 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
3727
3728 static PyMethodDef repeat_methods[] = {
3729 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3730 {NULL, NULL} /* sentinel */
3731 };
3732
3733 PyDoc_STRVAR(repeat_doc,
3734 "repeat(object [,times]) -> create an iterator which returns the object\n\
3735 for the specified number of times. If not specified, returns the object\n\
3736 endlessly.");
3737
3738 static PyTypeObject repeat_type = {
3739 PyVarObject_HEAD_INIT(NULL, 0)
3740 "itertools.repeat", /* tp_name */
3741 sizeof(repeatobject), /* tp_basicsize */
3742 0, /* tp_itemsize */
3743 /* methods */
3744 (destructor)repeat_dealloc, /* tp_dealloc */
3745 0, /* tp_print */
3746 0, /* tp_getattr */
3747 0, /* tp_setattr */
3748 0, /* tp_compare */
3749 (reprfunc)repeat_repr, /* tp_repr */
3750 0, /* tp_as_number */
3751 0, /* tp_as_sequence */
3752 0, /* tp_as_mapping */
3753 0, /* tp_hash */
3754 0, /* tp_call */
3755 0, /* tp_str */
3756 PyObject_GenericGetAttr, /* tp_getattro */
3757 0, /* tp_setattro */
3758 0, /* tp_as_buffer */
3759 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3760 Py_TPFLAGS_BASETYPE, /* tp_flags */
3761 repeat_doc, /* tp_doc */
3762 (traverseproc)repeat_traverse, /* tp_traverse */
3763 0, /* tp_clear */
3764 0, /* tp_richcompare */
3765 0, /* tp_weaklistoffset */
3766 PyObject_SelfIter, /* tp_iter */
3767 (iternextfunc)repeat_next, /* tp_iternext */
3768 repeat_methods, /* tp_methods */
3769 0, /* tp_members */
3770 0, /* tp_getset */
3771 0, /* tp_base */
3772 0, /* tp_dict */
3773 0, /* tp_descr_get */
3774 0, /* tp_descr_set */
3775 0, /* tp_dictoffset */
3776 0, /* tp_init */
3777 0, /* tp_alloc */
3778 repeat_new, /* tp_new */
3779 PyObject_GC_Del, /* tp_free */
3780 };
3781
3782 /* iziplongest object ************************************************************/
3783
3784 #include "Python.h"
3785
3786 typedef struct {
3787 PyObject_HEAD
3788 Py_ssize_t tuplesize;
3789 Py_ssize_t numactive;
3790 PyObject *ittuple; /* tuple of iterators */
3791 PyObject *result;
3792 PyObject *fillvalue;
3793 } iziplongestobject;
3794
3795 static PyTypeObject iziplongest_type;
3796
3797 static PyObject *
3798 izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3799 {
3800 iziplongestobject *lz;
3801 Py_ssize_t i;
3802 PyObject *ittuple; /* tuple of iterators */
3803 PyObject *result;
3804 PyObject *fillvalue = Py_None;
3805 Py_ssize_t tuplesize = PySequence_Length(args);
3806
3807 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3808 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3809 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3810 PyErr_SetString(PyExc_TypeError,
3811 "izip_longest() got an unexpected keyword argument");
3812 return NULL;
3813 }
3814 }
3815
3816 /* args must be a tuple */
3817 assert(PyTuple_Check(args));
3818
3819 /* obtain iterators */
3820 ittuple = PyTuple_New(tuplesize);
3821 if (ittuple == NULL)
3822 return NULL;
3823 for (i=0; i < tuplesize; ++i) {
3824 PyObject *item = PyTuple_GET_ITEM(args, i);
3825 PyObject *it = PyObject_GetIter(item);
3826 if (it == NULL) {
3827 if (PyErr_ExceptionMatches(PyExc_TypeError))
3828 PyErr_Format(PyExc_TypeError,
3829 "izip_longest argument #%zd must support iteration",
3830 i+1);
3831 Py_DECREF(ittuple);
3832 return NULL;
3833 }
3834 PyTuple_SET_ITEM(ittuple, i, it);
3835 }
3836
3837 /* create a result holder */
3838 result = PyTuple_New(tuplesize);
3839 if (result == NULL) {
3840 Py_DECREF(ittuple);
3841 return NULL;
3842 }
3843 for (i=0 ; i < tuplesize ; i++) {
3844 Py_INCREF(Py_None);
3845 PyTuple_SET_ITEM(result, i, Py_None);
3846 }
3847
3848 /* create iziplongestobject structure */
3849 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3850 if (lz == NULL) {
3851 Py_DECREF(ittuple);
3852 Py_DECREF(result);
3853 return NULL;
3854 }
3855 lz->ittuple = ittuple;
3856 lz->tuplesize = tuplesize;
3857 lz->numactive = tuplesize;
3858 lz->result = result;
3859 Py_INCREF(fillvalue);
3860 lz->fillvalue = fillvalue;
3861 return (PyObject *)lz;
3862 }
3863
3864 static void
3865 izip_longest_dealloc(iziplongestobject *lz)
3866 {
3867 PyObject_GC_UnTrack(lz);
3868 Py_XDECREF(lz->ittuple);
3869 Py_XDECREF(lz->result);
3870 Py_XDECREF(lz->fillvalue);
3871 Py_TYPE(lz)->tp_free(lz);
3872 }
3873
3874 static int
3875 izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3876 {
3877 Py_VISIT(lz->ittuple);
3878 Py_VISIT(lz->result);
3879 Py_VISIT(lz->fillvalue);
3880 return 0;
3881 }
3882
3883 static PyObject *
3884 izip_longest_next(iziplongestobject *lz)
3885 {
3886 Py_ssize_t i;
3887 Py_ssize_t tuplesize = lz->tuplesize;
3888 PyObject *result = lz->result;
3889 PyObject *it;
3890 PyObject *item;
3891 PyObject *olditem;
3892
3893 if (tuplesize == 0)
3894 return NULL;
3895 if (lz->numactive == 0)
3896 return NULL;
3897 if (Py_REFCNT(result) == 1) {
3898 Py_INCREF(result);
3899 for (i=0 ; i < tuplesize ; i++) {
3900 it = PyTuple_GET_ITEM(lz->ittuple, i);
3901 if (it == NULL) {
3902 Py_INCREF(lz->fillvalue);
3903 item = lz->fillvalue;
3904 } else {
3905 item = PyIter_Next(it);
3906 if (item == NULL) {
3907 lz->numactive -= 1;
3908 if (lz->numactive == 0 || PyErr_Occurred()) {
3909 lz->numactive = 0;
3910 Py_DECREF(result);
3911 return NULL;
3912 } else {
3913 Py_INCREF(lz->fillvalue);
3914 item = lz->fillvalue;
3915 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3916 Py_DECREF(it);
3917 }
3918 }
3919 }
3920 olditem = PyTuple_GET_ITEM(result, i);
3921 PyTuple_SET_ITEM(result, i, item);
3922 Py_DECREF(olditem);
3923 }
3924 } else {
3925 result = PyTuple_New(tuplesize);
3926 if (result == NULL)
3927 return NULL;
3928 for (i=0 ; i < tuplesize ; i++) {
3929 it = PyTuple_GET_ITEM(lz->ittuple, i);
3930 if (it == NULL) {
3931 Py_INCREF(lz->fillvalue);
3932 item = lz->fillvalue;
3933 } else {
3934 item = PyIter_Next(it);
3935 if (item == NULL) {
3936 lz->numactive -= 1;
3937 if (lz->numactive == 0 || PyErr_Occurred()) {
3938 lz->numactive = 0;
3939 Py_DECREF(result);
3940 return NULL;
3941 } else {
3942 Py_INCREF(lz->fillvalue);
3943 item = lz->fillvalue;
3944 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3945 Py_DECREF(it);
3946 }
3947 }
3948 }
3949 PyTuple_SET_ITEM(result, i, item);
3950 }
3951 }
3952 return result;
3953 }
3954
3955 PyDoc_STRVAR(izip_longest_doc,
3956 "izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3957 \n\
3958 Return an izip_longest object whose .next() method returns a tuple where\n\
3959 the i-th element comes from the i-th iterable argument. The .next()\n\
3960 method continues until the longest iterable in the argument sequence\n\
3961 is exhausted and then it raises StopIteration. When the shorter iterables\n\
3962 are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3963 defaults to None or can be specified by a keyword argument.\n\
3964 ");
3965
3966 static PyTypeObject iziplongest_type = {
3967 PyVarObject_HEAD_INIT(NULL, 0)
3968 "itertools.izip_longest", /* tp_name */
3969 sizeof(iziplongestobject), /* tp_basicsize */
3970 0, /* tp_itemsize */
3971 /* methods */
3972 (destructor)izip_longest_dealloc, /* tp_dealloc */
3973 0, /* tp_print */
3974 0, /* tp_getattr */
3975 0, /* tp_setattr */
3976 0, /* tp_compare */
3977 0, /* tp_repr */
3978 0, /* tp_as_number */
3979 0, /* tp_as_sequence */
3980 0, /* tp_as_mapping */
3981 0, /* tp_hash */
3982 0, /* tp_call */
3983 0, /* tp_str */
3984 PyObject_GenericGetAttr, /* tp_getattro */
3985 0, /* tp_setattro */
3986 0, /* tp_as_buffer */
3987 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3988 Py_TPFLAGS_BASETYPE, /* tp_flags */
3989 izip_longest_doc, /* tp_doc */
3990 (traverseproc)izip_longest_traverse, /* tp_traverse */
3991 0, /* tp_clear */
3992 0, /* tp_richcompare */
3993 0, /* tp_weaklistoffset */
3994 PyObject_SelfIter, /* tp_iter */
3995 (iternextfunc)izip_longest_next, /* tp_iternext */
3996 0, /* tp_methods */
3997 0, /* tp_members */
3998 0, /* tp_getset */
3999 0, /* tp_base */
4000 0, /* tp_dict */
4001 0, /* tp_descr_get */
4002 0, /* tp_descr_set */
4003 0, /* tp_dictoffset */
4004 0, /* tp_init */
4005 0, /* tp_alloc */
4006 izip_longest_new, /* tp_new */
4007 PyObject_GC_Del, /* tp_free */
4008 };
4009
4010 /* module level code ********************************************************/
4011
4012 PyDoc_STRVAR(module_doc,
4013 "Functional tools for creating and using iterators.\n\
4014 \n\
4015 Infinite iterators:\n\
4016 count([n]) --> n, n+1, n+2, ...\n\
4017 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
4018 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
4019 \n\
4020 Iterators terminating on the shortest input sequence:\n\
4021 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
4022 compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4023 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4024 groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
4025 ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4026 ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
4027 islice(seq, [start,] stop [, step]) --> elements from\n\
4028 seq[start:stop:step]\n\
4029 imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4030 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
4031 tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
4032 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
4033 izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4034 izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4035 \n\
4036 Combinatoric generators:\n\
4037 product(p, q, ... [repeat=1]) --> cartesian product\n\
4038 permutations(p[, r])\n\
4039 combinations(p, r)\n\
4040 combinations_with_replacement(p, r)\n\
4041 ");
4042
4043
4044 static PyMethodDef module_methods[] = {
4045 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4046 {NULL, NULL} /* sentinel */
4047 };
4048
4049 PyMODINIT_FUNC
4050 inititertools(void)
4051 {
4052 int i;
4053 PyObject *m;
4054 char *name;
4055 PyTypeObject *typelist[] = {
4056 &combinations_type,
4057 &cwr_type,
4058 &cycle_type,
4059 &dropwhile_type,
4060 &takewhile_type,
4061 &islice_type,
4062 &starmap_type,
4063 &imap_type,
4064 &chain_type,
4065 &compress_type,
4066 &ifilter_type,
4067 &ifilterfalse_type,
4068 &count_type,
4069 &izip_type,
4070 &iziplongest_type,
4071 &permutations_type,
4072 &product_type,
4073 &repeat_type,
4074 &groupby_type,
4075 NULL
4076 };
4077
4078 Py_TYPE(&teedataobject_type) = &PyType_Type;
4079 m = Py_InitModule3("itertools", module_methods, module_doc);
4080 if (m == NULL)
4081 return;
4082
4083 for (i=0 ; typelist[i] != NULL ; i++) {
4084 if (PyType_Ready(typelist[i]) < 0)
4085 return;
4086 name = strchr(typelist[i]->tp_name, '.');
4087 assert (name != NULL);
4088 Py_INCREF(typelist[i]);
4089 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4090 }
4091
4092 if (PyType_Ready(&teedataobject_type) < 0)
4093 return;
4094 if (PyType_Ready(&tee_type) < 0)
4095 return;
4096 if (PyType_Ready(&_grouper_type) < 0)
4097 return;
4098 }