only use -fno-strict-aliasing when needed by compiler
[python/dscho.git] / Modules / itertoolsmodule.c
blobba8d6dfde3d218b6c6af930c81b567b416abb9b1
2 #include "Python.h"
3 #include "structmember.h"
5 /* Itertools module written and maintained
6 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9 */
12 /* groupby object ***********************************************************/
14 typedef struct {
15 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
21 } groupbyobject;
23 static PyTypeObject groupby_type;
24 static PyObject *_grouper_create(groupbyobject *, PyObject *);
26 static PyObject *
27 groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
29 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
33 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
50 return (PyObject *)gbo;
53 static void
54 groupby_dealloc(groupbyobject *gbo)
56 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
65 static int
66 groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
68 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
76 static PyObject *
77 groupby_next(groupbyobject *gbo)
79 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
81 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
90 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
98 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
137 PyDoc_STRVAR(groupby_doc,
138 "groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139 (key, sub-iterator) grouped by each value of key(value).\n");
141 static PyTypeObject groupby_type = {
142 PyVarObject_HEAD_INIT(NULL, 0)
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_reserved */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
186 /* _grouper object (internal) ************************************************/
188 typedef struct {
189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
192 } _grouperobject;
194 static PyTypeObject _grouper_type;
196 static PyObject *
197 _grouper_create(groupbyobject *parent, PyObject *tgtkey)
199 _grouperobject *igo;
201 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
209 PyObject_GC_Track(igo);
210 return (PyObject *)igo;
213 static void
214 _grouper_dealloc(_grouperobject *igo)
216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
222 static int
223 _grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
230 static PyObject *
231 _grouper_next(_grouperobject *igo)
233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
242 if (gbo->keyfunc == Py_None) {
243 newkey = newvalue;
244 Py_INCREF(newvalue);
245 } else {
246 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
247 newvalue, NULL);
248 if (newkey == NULL) {
249 Py_DECREF(newvalue);
250 return NULL;
254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
259 assert(gbo->currkey != NULL);
260 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
261 if (rcmp <= 0)
262 /* got any error or current group is end */
263 return NULL;
265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
269 return r;
272 static PyTypeObject _grouper_type = {
273 PyVarObject_HEAD_INIT(NULL, 0)
274 "itertools._grouper", /* tp_name */
275 sizeof(_grouperobject), /* tp_basicsize */
276 0, /* tp_itemsize */
277 /* methods */
278 (destructor)_grouper_dealloc, /* tp_dealloc */
279 0, /* tp_print */
280 0, /* tp_getattr */
281 0, /* tp_setattr */
282 0, /* tp_reserved */
283 0, /* tp_repr */
284 0, /* tp_as_number */
285 0, /* tp_as_sequence */
286 0, /* tp_as_mapping */
287 0, /* tp_hash */
288 0, /* tp_call */
289 0, /* tp_str */
290 PyObject_GenericGetAttr, /* tp_getattro */
291 0, /* tp_setattro */
292 0, /* tp_as_buffer */
293 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
294 0, /* tp_doc */
295 (traverseproc)_grouper_traverse,/* tp_traverse */
296 0, /* tp_clear */
297 0, /* tp_richcompare */
298 0, /* tp_weaklistoffset */
299 PyObject_SelfIter, /* tp_iter */
300 (iternextfunc)_grouper_next, /* tp_iternext */
301 0, /* tp_methods */
302 0, /* tp_members */
303 0, /* tp_getset */
304 0, /* tp_base */
305 0, /* tp_dict */
306 0, /* tp_descr_get */
307 0, /* tp_descr_set */
308 0, /* tp_dictoffset */
309 0, /* tp_init */
310 0, /* tp_alloc */
311 0, /* tp_new */
312 PyObject_GC_Del, /* tp_free */
317 /* tee object and with supporting function and objects ***************/
319 /* The teedataobject pre-allocates space for LINKCELLS number of objects.
320 To help the object fit neatly inside cache lines (space for 16 to 32
321 pointers), the value should be a multiple of 16 minus space for
322 the other structure members including PyHEAD overhead. The larger the
323 value, the less memory overhead per object and the less time spent
324 allocating/deallocating new links. The smaller the number, the less
325 wasted space and the more rapid freeing of older data.
327 #define LINKCELLS 57
329 typedef struct {
330 PyObject_HEAD
331 PyObject *it;
332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
335 } teedataobject;
337 typedef struct {
338 PyObject_HEAD
339 teedataobject *dataobj;
340 int index;
341 PyObject *weakreflist;
342 } teeobject;
344 static PyTypeObject teedataobject_type;
346 static PyObject *
347 teedataobject_new(PyObject *it)
349 teedataobject *tdo;
351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
355 tdo->numread = 0;
356 tdo->nextlink = NULL;
357 Py_INCREF(it);
358 tdo->it = it;
359 PyObject_GC_Track(tdo);
360 return (PyObject *)tdo;
363 static PyObject *
364 teedataobject_jumplink(teedataobject *tdo)
366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
368 Py_XINCREF(tdo->nextlink);
369 return tdo->nextlink;
372 static PyObject *
373 teedataobject_getitem(teedataobject *tdo, int i)
375 PyObject *value;
377 assert(i < LINKCELLS);
378 if (i < tdo->numread)
379 value = tdo->values[i];
380 else {
381 /* this is the lead iterator, so fetch more data */
382 assert(i == tdo->numread);
383 value = PyIter_Next(tdo->it);
384 if (value == NULL)
385 return NULL;
386 tdo->numread++;
387 tdo->values[i] = value;
389 Py_INCREF(value);
390 return value;
393 static int
394 teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
396 int i;
397 Py_VISIT(tdo->it);
398 for (i = 0; i < tdo->numread; i++)
399 Py_VISIT(tdo->values[i]);
400 Py_VISIT(tdo->nextlink);
401 return 0;
404 static int
405 teedataobject_clear(teedataobject *tdo)
407 int i;
408 Py_CLEAR(tdo->it);
409 for (i=0 ; i<tdo->numread ; i++)
410 Py_CLEAR(tdo->values[i]);
411 Py_CLEAR(tdo->nextlink);
412 return 0;
415 static void
416 teedataobject_dealloc(teedataobject *tdo)
418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
423 PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
425 static PyTypeObject teedataobject_type = {
426 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
427 "itertools.tee_dataobject", /* tp_name */
428 sizeof(teedataobject), /* tp_basicsize */
429 0, /* tp_itemsize */
430 /* methods */
431 (destructor)teedataobject_dealloc, /* tp_dealloc */
432 0, /* tp_print */
433 0, /* tp_getattr */
434 0, /* tp_setattr */
435 0, /* tp_reserved */
436 0, /* tp_repr */
437 0, /* tp_as_number */
438 0, /* tp_as_sequence */
439 0, /* tp_as_mapping */
440 0, /* tp_hash */
441 0, /* tp_call */
442 0, /* tp_str */
443 PyObject_GenericGetAttr, /* tp_getattro */
444 0, /* tp_setattro */
445 0, /* tp_as_buffer */
446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
447 teedataobject_doc, /* tp_doc */
448 (traverseproc)teedataobject_traverse, /* tp_traverse */
449 (inquiry)teedataobject_clear, /* tp_clear */
450 0, /* tp_richcompare */
451 0, /* tp_weaklistoffset */
452 0, /* tp_iter */
453 0, /* tp_iternext */
454 0, /* tp_methods */
455 0, /* tp_members */
456 0, /* tp_getset */
457 0, /* tp_base */
458 0, /* tp_dict */
459 0, /* tp_descr_get */
460 0, /* tp_descr_set */
461 0, /* tp_dictoffset */
462 0, /* tp_init */
463 0, /* tp_alloc */
464 0, /* tp_new */
465 PyObject_GC_Del, /* tp_free */
469 static PyTypeObject tee_type;
471 static PyObject *
472 tee_next(teeobject *to)
474 PyObject *value, *link;
476 if (to->index >= LINKCELLS) {
477 link = teedataobject_jumplink(to->dataobj);
478 Py_DECREF(to->dataobj);
479 to->dataobj = (teedataobject *)link;
480 to->index = 0;
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
489 static int
490 tee_traverse(teeobject *to, visitproc visit, void *arg)
492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
496 static PyObject *
497 tee_copy(teeobject *to)
499 teeobject *newto;
501 newto = PyObject_GC_New(teeobject, &tee_type);
502 if (newto == NULL)
503 return NULL;
504 Py_INCREF(to->dataobj);
505 newto->dataobj = to->dataobj;
506 newto->index = to->index;
507 newto->weakreflist = NULL;
508 PyObject_GC_Track(newto);
509 return (PyObject *)newto;
512 PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
514 static PyObject *
515 tee_fromiterable(PyObject *iterable)
517 teeobject *to;
518 PyObject *it = NULL;
520 it = PyObject_GetIter(iterable);
521 if (it == NULL)
522 return NULL;
523 if (PyObject_TypeCheck(it, &tee_type)) {
524 to = (teeobject *)tee_copy((teeobject *)it);
525 goto done;
528 to = PyObject_GC_New(teeobject, &tee_type);
529 if (to == NULL)
530 goto done;
531 to->dataobj = (teedataobject *)teedataobject_new(it);
532 if (!to->dataobj) {
533 PyObject_GC_Del(to);
534 to = NULL;
535 goto done;
538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
541 done:
542 Py_XDECREF(it);
543 return (PyObject *)to;
546 static PyObject *
547 tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
549 PyObject *iterable;
551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
553 return tee_fromiterable(iterable);
556 static int
557 tee_clear(teeobject *to)
559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
561 Py_CLEAR(to->dataobj);
562 return 0;
565 static void
566 tee_dealloc(teeobject *to)
568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
573 PyDoc_STRVAR(teeobject_doc,
574 "Iterator wrapped to make it copyable");
576 static PyMethodDef tee_methods[] = {
577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
581 static PyTypeObject tee_type = {
582 PyVarObject_HEAD_INIT(NULL, 0)
583 "itertools.tee", /* tp_name */
584 sizeof(teeobject), /* tp_basicsize */
585 0, /* tp_itemsize */
586 /* methods */
587 (destructor)tee_dealloc, /* tp_dealloc */
588 0, /* tp_print */
589 0, /* tp_getattr */
590 0, /* tp_setattr */
591 0, /* tp_reserved */
592 0, /* tp_repr */
593 0, /* tp_as_number */
594 0, /* tp_as_sequence */
595 0, /* tp_as_mapping */
596 0, /* tp_hash */
597 0, /* tp_call */
598 0, /* tp_str */
599 0, /* tp_getattro */
600 0, /* tp_setattro */
601 0, /* tp_as_buffer */
602 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
603 teeobject_doc, /* tp_doc */
604 (traverseproc)tee_traverse, /* tp_traverse */
605 (inquiry)tee_clear, /* tp_clear */
606 0, /* tp_richcompare */
607 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
608 PyObject_SelfIter, /* tp_iter */
609 (iternextfunc)tee_next, /* tp_iternext */
610 tee_methods, /* tp_methods */
611 0, /* tp_members */
612 0, /* tp_getset */
613 0, /* tp_base */
614 0, /* tp_dict */
615 0, /* tp_descr_get */
616 0, /* tp_descr_set */
617 0, /* tp_dictoffset */
618 0, /* tp_init */
619 0, /* tp_alloc */
620 tee_new, /* tp_new */
621 PyObject_GC_Del, /* tp_free */
624 static PyObject *
625 tee(PyObject *self, PyObject *args)
627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
630 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
631 return NULL;
632 if (n < 0) {
633 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
634 return NULL;
636 result = PyTuple_New(n);
637 if (result == NULL)
638 return NULL;
639 if (n == 0)
640 return result;
641 it = PyObject_GetIter(iterable);
642 if (it == NULL) {
643 Py_DECREF(result);
644 return NULL;
646 if (!PyObject_HasAttrString(it, "__copy__")) {
647 copyable = tee_fromiterable(it);
648 Py_DECREF(it);
649 if (copyable == NULL) {
650 Py_DECREF(result);
651 return NULL;
653 } else
654 copyable = it;
655 PyTuple_SET_ITEM(result, 0, copyable);
656 for (i=1 ; i<n ; i++) {
657 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
658 if (copyable == NULL) {
659 Py_DECREF(result);
660 return NULL;
662 PyTuple_SET_ITEM(result, i, copyable);
664 return result;
667 PyDoc_STRVAR(tee_doc,
668 "tee(iterable, n=2) --> tuple of n independent iterators.");
671 /* cycle object **********************************************************/
673 typedef struct {
674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
678 } cycleobject;
680 static PyTypeObject cycle_type;
682 static PyObject *
683 cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
707 /* create cycleobject structure */
708 lz = (cycleobject *)type->tp_alloc(type, 0);
709 if (lz == NULL) {
710 Py_DECREF(it);
711 Py_DECREF(saved);
712 return NULL;
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
718 return (PyObject *)lz;
721 static void
722 cycle_dealloc(cycleobject *lz)
724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
727 Py_TYPE(lz)->tp_free(lz);
730 static int
731 cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
735 return 0;
738 static PyObject *
739 cycle_next(cycleobject *lz)
741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
752 return item;
754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
765 tmp = lz->it;
766 lz->it = it;
767 lz->firstpass = 1;
768 Py_DECREF(tmp);
772 PyDoc_STRVAR(cycle_doc,
773 "cycle(iterable) --> cycle object\n\
775 Return elements from the iterable until it is exhausted.\n\
776 Then repeat the sequence indefinitely.");
778 static PyTypeObject cycle_type = {
779 PyVarObject_HEAD_INIT(NULL, 0)
780 "itertools.cycle", /* tp_name */
781 sizeof(cycleobject), /* tp_basicsize */
782 0, /* tp_itemsize */
783 /* methods */
784 (destructor)cycle_dealloc, /* tp_dealloc */
785 0, /* tp_print */
786 0, /* tp_getattr */
787 0, /* tp_setattr */
788 0, /* tp_reserved */
789 0, /* tp_repr */
790 0, /* tp_as_number */
791 0, /* tp_as_sequence */
792 0, /* tp_as_mapping */
793 0, /* tp_hash */
794 0, /* tp_call */
795 0, /* tp_str */
796 PyObject_GenericGetAttr, /* tp_getattro */
797 0, /* tp_setattro */
798 0, /* tp_as_buffer */
799 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
800 Py_TPFLAGS_BASETYPE, /* tp_flags */
801 cycle_doc, /* tp_doc */
802 (traverseproc)cycle_traverse, /* tp_traverse */
803 0, /* tp_clear */
804 0, /* tp_richcompare */
805 0, /* tp_weaklistoffset */
806 PyObject_SelfIter, /* tp_iter */
807 (iternextfunc)cycle_next, /* tp_iternext */
808 0, /* tp_methods */
809 0, /* tp_members */
810 0, /* tp_getset */
811 0, /* tp_base */
812 0, /* tp_dict */
813 0, /* tp_descr_get */
814 0, /* tp_descr_set */
815 0, /* tp_dictoffset */
816 0, /* tp_init */
817 0, /* tp_alloc */
818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
823 /* dropwhile object **********************************************************/
825 typedef struct {
826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
830 } dropwhileobject;
832 static PyTypeObject dropwhile_type;
834 static PyObject *
835 dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
852 /* create dropwhileobject structure */
853 lz = (dropwhileobject *)type->tp_alloc(type, 0);
854 if (lz == NULL) {
855 Py_DECREF(it);
856 return NULL;
858 Py_INCREF(func);
859 lz->func = func;
860 lz->it = it;
861 lz->start = 0;
863 return (PyObject *)lz;
866 static void
867 dropwhile_dealloc(dropwhileobject *lz)
869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
872 Py_TYPE(lz)->tp_free(lz);
875 static int
876 dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
880 return 0;
883 static PyObject *
884 dropwhile_next(dropwhileobject *lz)
886 PyObject *item, *good;
887 PyObject *it = lz->it;
888 long ok;
889 PyObject *(*iternext)(PyObject *);
891 iternext = *Py_TYPE(it)->tp_iternext;
892 for (;;) {
893 item = iternext(it);
894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
910 Py_DECREF(item);
914 PyDoc_STRVAR(dropwhile_doc,
915 "dropwhile(predicate, iterable) --> dropwhile object\n\
917 Drop items from the iterable while predicate(item) is true.\n\
918 Afterwards, return every element until the iterable is exhausted.");
920 static PyTypeObject dropwhile_type = {
921 PyVarObject_HEAD_INIT(NULL, 0)
922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
924 0, /* tp_itemsize */
925 /* methods */
926 (destructor)dropwhile_dealloc, /* tp_dealloc */
927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_reserved */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
943 dropwhile_doc, /* tp_doc */
944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
948 PyObject_SelfIter, /* tp_iter */
949 (iternextfunc)dropwhile_next, /* tp_iternext */
950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
959 0, /* tp_alloc */
960 dropwhile_new, /* tp_new */
961 PyObject_GC_Del, /* tp_free */
965 /* takewhile object **********************************************************/
967 typedef struct {
968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
972 } takewhileobject;
974 static PyTypeObject takewhile_type;
976 static PyObject *
977 takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
1005 return (PyObject *)lz;
1008 static void
1009 takewhile_dealloc(takewhileobject *lz)
1011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
1014 Py_TYPE(lz)->tp_free(lz);
1017 static int
1018 takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
1022 return 0;
1025 static PyObject *
1026 takewhile_next(takewhileobject *lz)
1028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
1032 if (lz->stop == 1)
1033 return NULL;
1035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
1039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1044 ok = PyObject_IsTrue(good);
1045 Py_DECREF(good);
1046 if (ok)
1047 return item;
1048 Py_DECREF(item);
1049 lz->stop = 1;
1050 return NULL;
1053 PyDoc_STRVAR(takewhile_doc,
1054 "takewhile(predicate, iterable) --> takewhile object\n\
1056 Return successive entries from an iterable as long as the \n\
1057 predicate evaluates to true for each entry.");
1059 static PyTypeObject takewhile_type = {
1060 PyVarObject_HEAD_INIT(NULL, 0)
1061 "itertools.takewhile", /* tp_name */
1062 sizeof(takewhileobject), /* tp_basicsize */
1063 0, /* tp_itemsize */
1064 /* methods */
1065 (destructor)takewhile_dealloc, /* tp_dealloc */
1066 0, /* tp_print */
1067 0, /* tp_getattr */
1068 0, /* tp_setattr */
1069 0, /* tp_reserved */
1070 0, /* tp_repr */
1071 0, /* tp_as_number */
1072 0, /* tp_as_sequence */
1073 0, /* tp_as_mapping */
1074 0, /* tp_hash */
1075 0, /* tp_call */
1076 0, /* tp_str */
1077 PyObject_GenericGetAttr, /* tp_getattro */
1078 0, /* tp_setattro */
1079 0, /* tp_as_buffer */
1080 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1081 Py_TPFLAGS_BASETYPE, /* tp_flags */
1082 takewhile_doc, /* tp_doc */
1083 (traverseproc)takewhile_traverse, /* tp_traverse */
1084 0, /* tp_clear */
1085 0, /* tp_richcompare */
1086 0, /* tp_weaklistoffset */
1087 PyObject_SelfIter, /* tp_iter */
1088 (iternextfunc)takewhile_next, /* tp_iternext */
1089 0, /* tp_methods */
1090 0, /* tp_members */
1091 0, /* tp_getset */
1092 0, /* tp_base */
1093 0, /* tp_dict */
1094 0, /* tp_descr_get */
1095 0, /* tp_descr_set */
1096 0, /* tp_dictoffset */
1097 0, /* tp_init */
1098 0, /* tp_alloc */
1099 takewhile_new, /* tp_new */
1100 PyObject_GC_Del, /* tp_free */
1104 /* islice object ************************************************************/
1106 typedef struct {
1107 PyObject_HEAD
1108 PyObject *it;
1109 Py_ssize_t next;
1110 Py_ssize_t stop;
1111 Py_ssize_t step;
1112 Py_ssize_t cnt;
1113 } isliceobject;
1115 static PyTypeObject islice_type;
1117 static PyObject *
1118 islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1120 PyObject *seq;
1121 Py_ssize_t start=0, stop=-1, step=1;
1122 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1123 Py_ssize_t numargs;
1124 isliceobject *lz;
1126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
1129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
1132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyLong_AsSsize_t(a1);
1136 if (stop == -1) {
1137 if (PyErr_Occurred())
1138 PyErr_Clear();
1139 PyErr_SetString(PyExc_ValueError,
1140 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1141 return NULL;
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyLong_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyLong_AsSsize_t(a2);
1151 if (stop == -1) {
1152 if (PyErr_Occurred())
1153 PyErr_Clear();
1154 PyErr_SetString(PyExc_ValueError,
1155 "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1156 return NULL;
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.");
1163 return NULL;
1166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyLong_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
1183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
1195 return (PyObject *)lz;
1198 static void
1199 islice_dealloc(isliceobject *lz)
1201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
1206 static int
1207 islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1209 Py_VISIT(lz->it);
1210 return 0;
1213 static PyObject *
1214 islice_next(isliceobject *lz)
1216 PyObject *item;
1217 PyObject *it = lz->it;
1218 Py_ssize_t oldnext;
1219 PyObject *(*iternext)(PyObject *);
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++;
1229 if (lz->stop != -1 && lz->cnt >= lz->stop)
1230 return NULL;
1231 item = iternext(it);
1232 if (item == NULL)
1233 return NULL;
1234 lz->cnt++;
1235 oldnext = lz->next;
1236 lz->next += lz->step;
1237 if (lz->next < oldnext) /* Check for overflow */
1238 lz->next = lz->stop;
1239 return item;
1242 PyDoc_STRVAR(islice_doc,
1243 "islice(iterable, [start,] stop [, step]) --> islice object\n\
1245 Return an iterator whose next() method returns selected values from an\n\
1246 iterable. If start is specified, will skip all preceding elements;\n\
1247 otherwise, start defaults to zero. Step defaults to one. If\n\
1248 specified as another value, step determines how many values are \n\
1249 skipped between successive calls. Works like a slice() on a list\n\
1250 but returns an iterator.");
1252 static PyTypeObject islice_type = {
1253 PyVarObject_HEAD_INIT(NULL, 0)
1254 "itertools.islice", /* tp_name */
1255 sizeof(isliceobject), /* tp_basicsize */
1256 0, /* tp_itemsize */
1257 /* methods */
1258 (destructor)islice_dealloc, /* tp_dealloc */
1259 0, /* tp_print */
1260 0, /* tp_getattr */
1261 0, /* tp_setattr */
1262 0, /* tp_reserved */
1263 0, /* tp_repr */
1264 0, /* tp_as_number */
1265 0, /* tp_as_sequence */
1266 0, /* tp_as_mapping */
1267 0, /* tp_hash */
1268 0, /* tp_call */
1269 0, /* tp_str */
1270 PyObject_GenericGetAttr, /* tp_getattro */
1271 0, /* tp_setattro */
1272 0, /* tp_as_buffer */
1273 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1274 Py_TPFLAGS_BASETYPE, /* tp_flags */
1275 islice_doc, /* tp_doc */
1276 (traverseproc)islice_traverse, /* tp_traverse */
1277 0, /* tp_clear */
1278 0, /* tp_richcompare */
1279 0, /* tp_weaklistoffset */
1280 PyObject_SelfIter, /* tp_iter */
1281 (iternextfunc)islice_next, /* tp_iternext */
1282 0, /* tp_methods */
1283 0, /* tp_members */
1284 0, /* tp_getset */
1285 0, /* tp_base */
1286 0, /* tp_dict */
1287 0, /* tp_descr_get */
1288 0, /* tp_descr_set */
1289 0, /* tp_dictoffset */
1290 0, /* tp_init */
1291 0, /* tp_alloc */
1292 islice_new, /* tp_new */
1293 PyObject_GC_Del, /* tp_free */
1297 /* starmap object ************************************************************/
1299 typedef struct {
1300 PyObject_HEAD
1301 PyObject *func;
1302 PyObject *it;
1303 } starmapobject;
1305 static PyTypeObject starmap_type;
1307 static PyObject *
1308 starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1310 PyObject *func, *seq;
1311 PyObject *it;
1312 starmapobject *lz;
1314 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1315 return NULL;
1317 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1318 return NULL;
1320 /* Get iterator. */
1321 it = PyObject_GetIter(seq);
1322 if (it == NULL)
1323 return NULL;
1325 /* create starmapobject structure */
1326 lz = (starmapobject *)type->tp_alloc(type, 0);
1327 if (lz == NULL) {
1328 Py_DECREF(it);
1329 return NULL;
1331 Py_INCREF(func);
1332 lz->func = func;
1333 lz->it = it;
1335 return (PyObject *)lz;
1338 static void
1339 starmap_dealloc(starmapobject *lz)
1341 PyObject_GC_UnTrack(lz);
1342 Py_XDECREF(lz->func);
1343 Py_XDECREF(lz->it);
1344 Py_TYPE(lz)->tp_free(lz);
1347 static int
1348 starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1350 Py_VISIT(lz->it);
1351 Py_VISIT(lz->func);
1352 return 0;
1355 static PyObject *
1356 starmap_next(starmapobject *lz)
1358 PyObject *args;
1359 PyObject *result;
1360 PyObject *it = lz->it;
1362 args = (*Py_TYPE(it)->tp_iternext)(it);
1363 if (args == NULL)
1364 return NULL;
1365 if (!PyTuple_CheckExact(args)) {
1366 PyObject *newargs = PySequence_Tuple(args);
1367 Py_DECREF(args);
1368 if (newargs == NULL)
1369 return NULL;
1370 args = newargs;
1372 result = PyObject_Call(lz->func, args, NULL);
1373 Py_DECREF(args);
1374 return result;
1377 PyDoc_STRVAR(starmap_doc,
1378 "starmap(function, sequence) --> starmap object\n\
1380 Return an iterator whose values are returned from the function evaluated\n\
1381 with a argument tuple taken from the given sequence.");
1383 static PyTypeObject starmap_type = {
1384 PyVarObject_HEAD_INIT(NULL, 0)
1385 "itertools.starmap", /* tp_name */
1386 sizeof(starmapobject), /* tp_basicsize */
1387 0, /* tp_itemsize */
1388 /* methods */
1389 (destructor)starmap_dealloc, /* tp_dealloc */
1390 0, /* tp_print */
1391 0, /* tp_getattr */
1392 0, /* tp_setattr */
1393 0, /* tp_reserved */
1394 0, /* tp_repr */
1395 0, /* tp_as_number */
1396 0, /* tp_as_sequence */
1397 0, /* tp_as_mapping */
1398 0, /* tp_hash */
1399 0, /* tp_call */
1400 0, /* tp_str */
1401 PyObject_GenericGetAttr, /* tp_getattro */
1402 0, /* tp_setattro */
1403 0, /* tp_as_buffer */
1404 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1405 Py_TPFLAGS_BASETYPE, /* tp_flags */
1406 starmap_doc, /* tp_doc */
1407 (traverseproc)starmap_traverse, /* tp_traverse */
1408 0, /* tp_clear */
1409 0, /* tp_richcompare */
1410 0, /* tp_weaklistoffset */
1411 PyObject_SelfIter, /* tp_iter */
1412 (iternextfunc)starmap_next, /* tp_iternext */
1413 0, /* tp_methods */
1414 0, /* tp_members */
1415 0, /* tp_getset */
1416 0, /* tp_base */
1417 0, /* tp_dict */
1418 0, /* tp_descr_get */
1419 0, /* tp_descr_set */
1420 0, /* tp_dictoffset */
1421 0, /* tp_init */
1422 0, /* tp_alloc */
1423 starmap_new, /* tp_new */
1424 PyObject_GC_Del, /* tp_free */
1428 /* chain object ************************************************************/
1430 typedef struct {
1431 PyObject_HEAD
1432 PyObject *source; /* Iterator over input iterables */
1433 PyObject *active; /* Currently running input iterator */
1434 } chainobject;
1436 static PyTypeObject chain_type;
1438 static PyObject *
1439 chain_new_internal(PyTypeObject *type, PyObject *source)
1441 chainobject *lz;
1443 lz = (chainobject *)type->tp_alloc(type, 0);
1444 if (lz == NULL) {
1445 Py_DECREF(source);
1446 return NULL;
1449 lz->source = source;
1450 lz->active = NULL;
1451 return (PyObject *)lz;
1454 static PyObject *
1455 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1457 PyObject *source;
1459 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1460 return NULL;
1462 source = PyObject_GetIter(args);
1463 if (source == NULL)
1464 return NULL;
1466 return chain_new_internal(type, source);
1469 static PyObject *
1470 chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1472 PyObject *source;
1474 source = PyObject_GetIter(arg);
1475 if (source == NULL)
1476 return NULL;
1478 return chain_new_internal(type, source);
1481 static void
1482 chain_dealloc(chainobject *lz)
1484 PyObject_GC_UnTrack(lz);
1485 Py_XDECREF(lz->active);
1486 Py_XDECREF(lz->source);
1487 Py_TYPE(lz)->tp_free(lz);
1490 static int
1491 chain_traverse(chainobject *lz, visitproc visit, void *arg)
1493 Py_VISIT(lz->source);
1494 Py_VISIT(lz->active);
1495 return 0;
1498 static PyObject *
1499 chain_next(chainobject *lz)
1501 PyObject *item;
1503 if (lz->source == NULL)
1504 return NULL; /* already stopped */
1506 if (lz->active == NULL) {
1507 PyObject *iterable = PyIter_Next(lz->source);
1508 if (iterable == NULL) {
1509 Py_CLEAR(lz->source);
1510 return NULL; /* no more input sources */
1512 lz->active = PyObject_GetIter(iterable);
1513 Py_DECREF(iterable);
1514 if (lz->active == NULL) {
1515 Py_CLEAR(lz->source);
1516 return NULL; /* input not iterable */
1519 item = PyIter_Next(lz->active);
1520 if (item != NULL)
1521 return item;
1522 if (PyErr_Occurred()) {
1523 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1524 PyErr_Clear();
1525 else
1526 return NULL; /* input raised an exception */
1528 Py_CLEAR(lz->active);
1529 return chain_next(lz); /* recurse and use next active */
1532 PyDoc_STRVAR(chain_doc,
1533 "chain(*iterables) --> chain object\n\
1535 Return a chain object whose .__next__() method returns elements from the\n\
1536 first iterable until it is exhausted, then elements from the next\n\
1537 iterable, until all of the iterables are exhausted.");
1539 PyDoc_STRVAR(chain_from_iterable_doc,
1540 "chain.from_iterable(iterable) --> chain object\n\
1542 Alternate chain() contructor taking a single iterable argument\n\
1543 that evaluates lazily.");
1545 static PyMethodDef chain_methods[] = {
1546 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1547 chain_from_iterable_doc},
1548 {NULL, NULL} /* sentinel */
1551 static PyTypeObject chain_type = {
1552 PyVarObject_HEAD_INIT(NULL, 0)
1553 "itertools.chain", /* tp_name */
1554 sizeof(chainobject), /* tp_basicsize */
1555 0, /* tp_itemsize */
1556 /* methods */
1557 (destructor)chain_dealloc, /* tp_dealloc */
1558 0, /* tp_print */
1559 0, /* tp_getattr */
1560 0, /* tp_setattr */
1561 0, /* tp_reserved */
1562 0, /* tp_repr */
1563 0, /* tp_as_number */
1564 0, /* tp_as_sequence */
1565 0, /* tp_as_mapping */
1566 0, /* tp_hash */
1567 0, /* tp_call */
1568 0, /* tp_str */
1569 PyObject_GenericGetAttr, /* tp_getattro */
1570 0, /* tp_setattro */
1571 0, /* tp_as_buffer */
1572 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1573 Py_TPFLAGS_BASETYPE, /* tp_flags */
1574 chain_doc, /* tp_doc */
1575 (traverseproc)chain_traverse, /* tp_traverse */
1576 0, /* tp_clear */
1577 0, /* tp_richcompare */
1578 0, /* tp_weaklistoffset */
1579 PyObject_SelfIter, /* tp_iter */
1580 (iternextfunc)chain_next, /* tp_iternext */
1581 chain_methods, /* tp_methods */
1582 0, /* tp_members */
1583 0, /* tp_getset */
1584 0, /* tp_base */
1585 0, /* tp_dict */
1586 0, /* tp_descr_get */
1587 0, /* tp_descr_set */
1588 0, /* tp_dictoffset */
1589 0, /* tp_init */
1590 0, /* tp_alloc */
1591 chain_new, /* tp_new */
1592 PyObject_GC_Del, /* tp_free */
1596 /* product object ************************************************************/
1598 typedef struct {
1599 PyObject_HEAD
1600 PyObject *pools; /* tuple of pool tuples */
1601 Py_ssize_t *indices; /* one index per pool */
1602 PyObject *result; /* most recently returned result tuple */
1603 int stopped; /* set to 1 when the product iterator is exhausted */
1604 } productobject;
1606 static PyTypeObject product_type;
1608 static PyObject *
1609 product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1611 productobject *lz;
1612 Py_ssize_t nargs, npools, repeat=1;
1613 PyObject *pools = NULL;
1614 Py_ssize_t *indices = NULL;
1615 Py_ssize_t i;
1617 if (kwds != NULL) {
1618 char *kwlist[] = {"repeat", 0};
1619 PyObject *tmpargs = PyTuple_New(0);
1620 if (tmpargs == NULL)
1621 return NULL;
1622 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1623 Py_DECREF(tmpargs);
1624 return NULL;
1626 Py_DECREF(tmpargs);
1627 if (repeat < 0) {
1628 PyErr_SetString(PyExc_ValueError,
1629 "repeat argument cannot be negative");
1630 return NULL;
1634 assert(PyTuple_Check(args));
1635 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1636 npools = nargs * repeat;
1638 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1639 if (indices == NULL) {
1640 PyErr_NoMemory();
1641 goto error;
1644 pools = PyTuple_New(npools);
1645 if (pools == NULL)
1646 goto error;
1648 for (i=0; i < nargs ; ++i) {
1649 PyObject *item = PyTuple_GET_ITEM(args, i);
1650 PyObject *pool = PySequence_Tuple(item);
1651 if (pool == NULL)
1652 goto error;
1653 PyTuple_SET_ITEM(pools, i, pool);
1654 indices[i] = 0;
1656 for ( ; i < npools; ++i) {
1657 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1658 Py_INCREF(pool);
1659 PyTuple_SET_ITEM(pools, i, pool);
1660 indices[i] = 0;
1663 /* create productobject structure */
1664 lz = (productobject *)type->tp_alloc(type, 0);
1665 if (lz == NULL)
1666 goto error;
1668 lz->pools = pools;
1669 lz->indices = indices;
1670 lz->result = NULL;
1671 lz->stopped = 0;
1673 return (PyObject *)lz;
1675 error:
1676 if (indices != NULL)
1677 PyMem_Free(indices);
1678 Py_XDECREF(pools);
1679 return NULL;
1682 static void
1683 product_dealloc(productobject *lz)
1685 PyObject_GC_UnTrack(lz);
1686 Py_XDECREF(lz->pools);
1687 Py_XDECREF(lz->result);
1688 if (lz->indices != NULL)
1689 PyMem_Free(lz->indices);
1690 Py_TYPE(lz)->tp_free(lz);
1693 static int
1694 product_traverse(productobject *lz, visitproc visit, void *arg)
1696 Py_VISIT(lz->pools);
1697 Py_VISIT(lz->result);
1698 return 0;
1701 static PyObject *
1702 product_next(productobject *lz)
1704 PyObject *pool;
1705 PyObject *elem;
1706 PyObject *oldelem;
1707 PyObject *pools = lz->pools;
1708 PyObject *result = lz->result;
1709 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1710 Py_ssize_t i;
1712 if (lz->stopped)
1713 return NULL;
1715 if (result == NULL) {
1716 /* On the first pass, return an initial tuple filled with the
1717 first element from each pool. */
1718 result = PyTuple_New(npools);
1719 if (result == NULL)
1720 goto empty;
1721 lz->result = result;
1722 for (i=0; i < npools; i++) {
1723 pool = PyTuple_GET_ITEM(pools, i);
1724 if (PyTuple_GET_SIZE(pool) == 0)
1725 goto empty;
1726 elem = PyTuple_GET_ITEM(pool, 0);
1727 Py_INCREF(elem);
1728 PyTuple_SET_ITEM(result, i, elem);
1730 } else {
1731 Py_ssize_t *indices = lz->indices;
1733 /* Copy the previous result tuple or re-use it if available */
1734 if (Py_REFCNT(result) > 1) {
1735 PyObject *old_result = result;
1736 result = PyTuple_New(npools);
1737 if (result == NULL)
1738 goto empty;
1739 lz->result = result;
1740 for (i=0; i < npools; i++) {
1741 elem = PyTuple_GET_ITEM(old_result, i);
1742 Py_INCREF(elem);
1743 PyTuple_SET_ITEM(result, i, elem);
1745 Py_DECREF(old_result);
1747 /* Now, we've got the only copy so we can update it in-place */
1748 assert (npools==0 || Py_REFCNT(result) == 1);
1750 /* Update the pool indices right-to-left. Only advance to the
1751 next pool when the previous one rolls-over */
1752 for (i=npools-1 ; i >= 0 ; i--) {
1753 pool = PyTuple_GET_ITEM(pools, i);
1754 indices[i]++;
1755 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1756 /* Roll-over and advance to next pool */
1757 indices[i] = 0;
1758 elem = PyTuple_GET_ITEM(pool, 0);
1759 Py_INCREF(elem);
1760 oldelem = PyTuple_GET_ITEM(result, i);
1761 PyTuple_SET_ITEM(result, i, elem);
1762 Py_DECREF(oldelem);
1763 } else {
1764 /* No rollover. Just increment and stop here. */
1765 elem = PyTuple_GET_ITEM(pool, indices[i]);
1766 Py_INCREF(elem);
1767 oldelem = PyTuple_GET_ITEM(result, i);
1768 PyTuple_SET_ITEM(result, i, elem);
1769 Py_DECREF(oldelem);
1770 break;
1774 /* If i is negative, then the indices have all rolled-over
1775 and we're done. */
1776 if (i < 0)
1777 goto empty;
1780 Py_INCREF(result);
1781 return result;
1783 empty:
1784 lz->stopped = 1;
1785 return NULL;
1788 PyDoc_STRVAR(product_doc,
1789 "product(*iterables) --> product object\n\
1791 Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
1792 For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1793 The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1794 cycle in a manner similar to an odometer (with the rightmost element changing\n\
1795 on every iteration).\n\n\
1796 To compute the product of an iterable with itself, specify the number\n\
1797 of repetitions with the optional repeat keyword argument. For example,\n\
1798 product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
1799 product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1800 product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1802 static PyTypeObject product_type = {
1803 PyVarObject_HEAD_INIT(NULL, 0)
1804 "itertools.product", /* tp_name */
1805 sizeof(productobject), /* tp_basicsize */
1806 0, /* tp_itemsize */
1807 /* methods */
1808 (destructor)product_dealloc, /* tp_dealloc */
1809 0, /* tp_print */
1810 0, /* tp_getattr */
1811 0, /* tp_setattr */
1812 0, /* tp_reserved */
1813 0, /* tp_repr */
1814 0, /* tp_as_number */
1815 0, /* tp_as_sequence */
1816 0, /* tp_as_mapping */
1817 0, /* tp_hash */
1818 0, /* tp_call */
1819 0, /* tp_str */
1820 PyObject_GenericGetAttr, /* tp_getattro */
1821 0, /* tp_setattro */
1822 0, /* tp_as_buffer */
1823 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1824 Py_TPFLAGS_BASETYPE, /* tp_flags */
1825 product_doc, /* tp_doc */
1826 (traverseproc)product_traverse, /* tp_traverse */
1827 0, /* tp_clear */
1828 0, /* tp_richcompare */
1829 0, /* tp_weaklistoffset */
1830 PyObject_SelfIter, /* tp_iter */
1831 (iternextfunc)product_next, /* tp_iternext */
1832 0, /* tp_methods */
1833 0, /* tp_members */
1834 0, /* tp_getset */
1835 0, /* tp_base */
1836 0, /* tp_dict */
1837 0, /* tp_descr_get */
1838 0, /* tp_descr_set */
1839 0, /* tp_dictoffset */
1840 0, /* tp_init */
1841 0, /* tp_alloc */
1842 product_new, /* tp_new */
1843 PyObject_GC_Del, /* tp_free */
1847 /* combinations object ************************************************************/
1849 typedef struct {
1850 PyObject_HEAD
1851 PyObject *pool; /* input converted to a tuple */
1852 Py_ssize_t *indices; /* one index per result element */
1853 PyObject *result; /* most recently returned result tuple */
1854 Py_ssize_t r; /* size of result tuple */
1855 int stopped; /* set to 1 when the combinations iterator is exhausted */
1856 } combinationsobject;
1858 static PyTypeObject combinations_type;
1860 static PyObject *
1861 combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1863 combinationsobject *co;
1864 Py_ssize_t n;
1865 Py_ssize_t r;
1866 PyObject *pool = NULL;
1867 PyObject *iterable = NULL;
1868 Py_ssize_t *indices = NULL;
1869 Py_ssize_t i;
1870 static char *kwargs[] = {"iterable", "r", NULL};
1872 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
1873 &iterable, &r))
1874 return NULL;
1876 pool = PySequence_Tuple(iterable);
1877 if (pool == NULL)
1878 goto error;
1879 n = PyTuple_GET_SIZE(pool);
1880 if (r < 0) {
1881 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
1882 goto error;
1885 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
1886 if (indices == NULL) {
1887 PyErr_NoMemory();
1888 goto error;
1891 for (i=0 ; i<r ; i++)
1892 indices[i] = i;
1894 /* create combinationsobject structure */
1895 co = (combinationsobject *)type->tp_alloc(type, 0);
1896 if (co == NULL)
1897 goto error;
1899 co->pool = pool;
1900 co->indices = indices;
1901 co->result = NULL;
1902 co->r = r;
1903 co->stopped = r > n ? 1 : 0;
1905 return (PyObject *)co;
1907 error:
1908 if (indices != NULL)
1909 PyMem_Free(indices);
1910 Py_XDECREF(pool);
1911 return NULL;
1914 static void
1915 combinations_dealloc(combinationsobject *co)
1917 PyObject_GC_UnTrack(co);
1918 Py_XDECREF(co->pool);
1919 Py_XDECREF(co->result);
1920 if (co->indices != NULL)
1921 PyMem_Free(co->indices);
1922 Py_TYPE(co)->tp_free(co);
1925 static int
1926 combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
1928 Py_VISIT(co->pool);
1929 Py_VISIT(co->result);
1930 return 0;
1933 static PyObject *
1934 combinations_next(combinationsobject *co)
1936 PyObject *elem;
1937 PyObject *oldelem;
1938 PyObject *pool = co->pool;
1939 Py_ssize_t *indices = co->indices;
1940 PyObject *result = co->result;
1941 Py_ssize_t n = PyTuple_GET_SIZE(pool);
1942 Py_ssize_t r = co->r;
1943 Py_ssize_t i, j, index;
1945 if (co->stopped)
1946 return NULL;
1948 if (result == NULL) {
1949 /* On the first pass, initialize result tuple using the indices */
1950 result = PyTuple_New(r);
1951 if (result == NULL)
1952 goto empty;
1953 co->result = result;
1954 for (i=0; i<r ; i++) {
1955 index = indices[i];
1956 elem = PyTuple_GET_ITEM(pool, index);
1957 Py_INCREF(elem);
1958 PyTuple_SET_ITEM(result, i, elem);
1960 } else {
1961 /* Copy the previous result tuple or re-use it if available */
1962 if (Py_REFCNT(result) > 1) {
1963 PyObject *old_result = result;
1964 result = PyTuple_New(r);
1965 if (result == NULL)
1966 goto empty;
1967 co->result = result;
1968 for (i=0; i<r ; i++) {
1969 elem = PyTuple_GET_ITEM(old_result, i);
1970 Py_INCREF(elem);
1971 PyTuple_SET_ITEM(result, i, elem);
1973 Py_DECREF(old_result);
1975 /* Now, we've got the only copy so we can update it in-place
1976 * CPython's empty tuple is a singleton and cached in
1977 * PyTuple's freelist.
1979 assert(r == 0 || Py_REFCNT(result) == 1);
1981 /* Scan indices right-to-left until finding one that is not
1982 at its maximum (i + n - r). */
1983 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
1986 /* If i is negative, then the indices are all at
1987 their maximum value and we're done. */
1988 if (i < 0)
1989 goto empty;
1991 /* Increment the current index which we know is not at its
1992 maximum. Then move back to the right setting each index
1993 to its lowest possible value (one higher than the index
1994 to its left -- this maintains the sort order invariant). */
1995 indices[i]++;
1996 for (j=i+1 ; j<r ; j++)
1997 indices[j] = indices[j-1] + 1;
1999 /* Update the result tuple for the new indices
2000 starting with i, the leftmost index that changed */
2001 for ( ; i<r ; i++) {
2002 index = indices[i];
2003 elem = PyTuple_GET_ITEM(pool, index);
2004 Py_INCREF(elem);
2005 oldelem = PyTuple_GET_ITEM(result, i);
2006 PyTuple_SET_ITEM(result, i, elem);
2007 Py_DECREF(oldelem);
2011 Py_INCREF(result);
2012 return result;
2014 empty:
2015 co->stopped = 1;
2016 return NULL;
2019 PyDoc_STRVAR(combinations_doc,
2020 "combinations(iterable, r) --> combinations object\n\
2022 Return successive r-length combinations of elements in the iterable.\n\n\
2023 combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2025 static PyTypeObject combinations_type = {
2026 PyVarObject_HEAD_INIT(NULL, 0)
2027 "itertools.combinations", /* tp_name */
2028 sizeof(combinationsobject), /* tp_basicsize */
2029 0, /* tp_itemsize */
2030 /* methods */
2031 (destructor)combinations_dealloc, /* tp_dealloc */
2032 0, /* tp_print */
2033 0, /* tp_getattr */
2034 0, /* tp_setattr */
2035 0, /* tp_reserved */
2036 0, /* tp_repr */
2037 0, /* tp_as_number */
2038 0, /* tp_as_sequence */
2039 0, /* tp_as_mapping */
2040 0, /* tp_hash */
2041 0, /* tp_call */
2042 0, /* tp_str */
2043 PyObject_GenericGetAttr, /* tp_getattro */
2044 0, /* tp_setattro */
2045 0, /* tp_as_buffer */
2046 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2047 Py_TPFLAGS_BASETYPE, /* tp_flags */
2048 combinations_doc, /* tp_doc */
2049 (traverseproc)combinations_traverse, /* tp_traverse */
2050 0, /* tp_clear */
2051 0, /* tp_richcompare */
2052 0, /* tp_weaklistoffset */
2053 PyObject_SelfIter, /* tp_iter */
2054 (iternextfunc)combinations_next, /* tp_iternext */
2055 0, /* tp_methods */
2056 0, /* tp_members */
2057 0, /* tp_getset */
2058 0, /* tp_base */
2059 0, /* tp_dict */
2060 0, /* tp_descr_get */
2061 0, /* tp_descr_set */
2062 0, /* tp_dictoffset */
2063 0, /* tp_init */
2064 0, /* tp_alloc */
2065 combinations_new, /* tp_new */
2066 PyObject_GC_Del, /* tp_free */
2070 /* combinations with replacement object *******************************************/
2072 /* Equivalent to:
2074 def combinations_with_replacement(iterable, r):
2075 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2076 # number items returned: (n+r-1)! / r! / (n-1)!
2077 pool = tuple(iterable)
2078 n = len(pool)
2079 indices = [0] * r
2080 yield tuple(pool[i] for i in indices)
2081 while 1:
2082 for i in reversed(range(r)):
2083 if indices[i] != n - 1:
2084 break
2085 else:
2086 return
2087 indices[i:] = [indices[i] + 1] * (r - i)
2088 yield tuple(pool[i] for i in indices)
2090 def combinations_with_replacement2(iterable, r):
2091 'Alternate version that filters from product()'
2092 pool = tuple(iterable)
2093 n = len(pool)
2094 for indices in product(range(n), repeat=r):
2095 if sorted(indices) == list(indices):
2096 yield tuple(pool[i] for i in indices)
2098 typedef struct {
2099 PyObject_HEAD
2100 PyObject *pool; /* input converted to a tuple */
2101 Py_ssize_t *indices; /* one index per result element */
2102 PyObject *result; /* most recently returned result tuple */
2103 Py_ssize_t r; /* size of result tuple */
2104 int stopped; /* set to 1 when the cwr iterator is exhausted */
2105 } cwrobject;
2107 static PyTypeObject cwr_type;
2109 static PyObject *
2110 cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2112 cwrobject *co;
2113 Py_ssize_t n;
2114 Py_ssize_t r;
2115 PyObject *pool = NULL;
2116 PyObject *iterable = NULL;
2117 Py_ssize_t *indices = NULL;
2118 Py_ssize_t i;
2119 static char *kwargs[] = {"iterable", "r", NULL};
2121 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2122 &iterable, &r))
2123 return NULL;
2125 pool = PySequence_Tuple(iterable);
2126 if (pool == NULL)
2127 goto error;
2128 n = PyTuple_GET_SIZE(pool);
2129 if (r < 0) {
2130 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2131 goto error;
2134 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2135 if (indices == NULL) {
2136 PyErr_NoMemory();
2137 goto error;
2140 for (i=0 ; i<r ; i++)
2141 indices[i] = 0;
2143 /* create cwrobject structure */
2144 co = (cwrobject *)type->tp_alloc(type, 0);
2145 if (co == NULL)
2146 goto error;
2148 co->pool = pool;
2149 co->indices = indices;
2150 co->result = NULL;
2151 co->r = r;
2152 co->stopped = !n && r;
2154 return (PyObject *)co;
2156 error:
2157 if (indices != NULL)
2158 PyMem_Free(indices);
2159 Py_XDECREF(pool);
2160 return NULL;
2163 static void
2164 cwr_dealloc(cwrobject *co)
2166 PyObject_GC_UnTrack(co);
2167 Py_XDECREF(co->pool);
2168 Py_XDECREF(co->result);
2169 if (co->indices != NULL)
2170 PyMem_Free(co->indices);
2171 Py_TYPE(co)->tp_free(co);
2174 static int
2175 cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2177 Py_VISIT(co->pool);
2178 Py_VISIT(co->result);
2179 return 0;
2182 static PyObject *
2183 cwr_next(cwrobject *co)
2185 PyObject *elem;
2186 PyObject *oldelem;
2187 PyObject *pool = co->pool;
2188 Py_ssize_t *indices = co->indices;
2189 PyObject *result = co->result;
2190 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2191 Py_ssize_t r = co->r;
2192 Py_ssize_t i, j, index;
2194 if (co->stopped)
2195 return NULL;
2197 if (result == NULL) {
2198 /* On the first pass, initialize result tuple using the indices */
2199 result = PyTuple_New(r);
2200 if (result == NULL)
2201 goto empty;
2202 co->result = result;
2203 for (i=0; i<r ; i++) {
2204 index = indices[i];
2205 elem = PyTuple_GET_ITEM(pool, index);
2206 Py_INCREF(elem);
2207 PyTuple_SET_ITEM(result, i, elem);
2209 } else {
2210 /* Copy the previous result tuple or re-use it if available */
2211 if (Py_REFCNT(result) > 1) {
2212 PyObject *old_result = result;
2213 result = PyTuple_New(r);
2214 if (result == NULL)
2215 goto empty;
2216 co->result = result;
2217 for (i=0; i<r ; i++) {
2218 elem = PyTuple_GET_ITEM(old_result, i);
2219 Py_INCREF(elem);
2220 PyTuple_SET_ITEM(result, i, elem);
2222 Py_DECREF(old_result);
2224 /* Now, we've got the only copy so we can update it in-place CPython's
2225 empty tuple is a singleton and cached in PyTuple's freelist. */
2226 assert(r == 0 || Py_REFCNT(result) == 1);
2228 /* Scan indices right-to-left until finding one that is not
2229 * at its maximum (n-1). */
2230 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2233 /* If i is negative, then the indices are all at
2234 their maximum value and we're done. */
2235 if (i < 0)
2236 goto empty;
2238 /* Increment the current index which we know is not at its
2239 maximum. Then set all to the right to the same value. */
2240 indices[i]++;
2241 for (j=i+1 ; j<r ; j++)
2242 indices[j] = indices[j-1];
2244 /* Update the result tuple for the new indices
2245 starting with i, the leftmost index that changed */
2246 for ( ; i<r ; i++) {
2247 index = indices[i];
2248 elem = PyTuple_GET_ITEM(pool, index);
2249 Py_INCREF(elem);
2250 oldelem = PyTuple_GET_ITEM(result, i);
2251 PyTuple_SET_ITEM(result, i, elem);
2252 Py_DECREF(oldelem);
2256 Py_INCREF(result);
2257 return result;
2259 empty:
2260 co->stopped = 1;
2261 return NULL;
2264 PyDoc_STRVAR(cwr_doc,
2265 "combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\
2267 Return successive r-length combinations of elements in the iterable\n\
2268 allowing individual elements to have successive repeats.\n\
2269 combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2271 static PyTypeObject cwr_type = {
2272 PyVarObject_HEAD_INIT(NULL, 0)
2273 "itertools.combinations_with_replacement", /* tp_name */
2274 sizeof(cwrobject), /* tp_basicsize */
2275 0, /* tp_itemsize */
2276 /* methods */
2277 (destructor)cwr_dealloc, /* tp_dealloc */
2278 0, /* tp_print */
2279 0, /* tp_getattr */
2280 0, /* tp_setattr */
2281 0, /* tp_reserved */
2282 0, /* tp_repr */
2283 0, /* tp_as_number */
2284 0, /* tp_as_sequence */
2285 0, /* tp_as_mapping */
2286 0, /* tp_hash */
2287 0, /* tp_call */
2288 0, /* tp_str */
2289 PyObject_GenericGetAttr, /* tp_getattro */
2290 0, /* tp_setattro */
2291 0, /* tp_as_buffer */
2292 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2293 Py_TPFLAGS_BASETYPE, /* tp_flags */
2294 cwr_doc, /* tp_doc */
2295 (traverseproc)cwr_traverse, /* tp_traverse */
2296 0, /* tp_clear */
2297 0, /* tp_richcompare */
2298 0, /* tp_weaklistoffset */
2299 PyObject_SelfIter, /* tp_iter */
2300 (iternextfunc)cwr_next, /* tp_iternext */
2301 0, /* tp_methods */
2302 0, /* tp_members */
2303 0, /* tp_getset */
2304 0, /* tp_base */
2305 0, /* tp_dict */
2306 0, /* tp_descr_get */
2307 0, /* tp_descr_set */
2308 0, /* tp_dictoffset */
2309 0, /* tp_init */
2310 0, /* tp_alloc */
2311 cwr_new, /* tp_new */
2312 PyObject_GC_Del, /* tp_free */
2316 /* permutations object ************************************************************
2318 def permutations(iterable, r=None):
2319 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2320 pool = tuple(iterable)
2321 n = len(pool)
2322 r = n if r is None else r
2323 indices = range(n)
2324 cycles = range(n-r+1, n+1)[::-1]
2325 yield tuple(pool[i] for i in indices[:r])
2326 while n:
2327 for i in reversed(range(r)):
2328 cycles[i] -= 1
2329 if cycles[i] == 0:
2330 indices[i:] = indices[i+1:] + indices[i:i+1]
2331 cycles[i] = n - i
2332 else:
2333 j = cycles[i]
2334 indices[i], indices[-j] = indices[-j], indices[i]
2335 yield tuple(pool[i] for i in indices[:r])
2336 break
2337 else:
2338 return
2341 typedef struct {
2342 PyObject_HEAD
2343 PyObject *pool; /* input converted to a tuple */
2344 Py_ssize_t *indices; /* one index per element in the pool */
2345 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2346 PyObject *result; /* most recently returned result tuple */
2347 Py_ssize_t r; /* size of result tuple */
2348 int stopped; /* set to 1 when the permutations iterator is exhausted */
2349 } permutationsobject;
2351 static PyTypeObject permutations_type;
2353 static PyObject *
2354 permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2356 permutationsobject *po;
2357 Py_ssize_t n;
2358 Py_ssize_t r;
2359 PyObject *robj = Py_None;
2360 PyObject *pool = NULL;
2361 PyObject *iterable = NULL;
2362 Py_ssize_t *indices = NULL;
2363 Py_ssize_t *cycles = NULL;
2364 Py_ssize_t i;
2365 static char *kwargs[] = {"iterable", "r", NULL};
2367 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2368 &iterable, &robj))
2369 return NULL;
2371 pool = PySequence_Tuple(iterable);
2372 if (pool == NULL)
2373 goto error;
2374 n = PyTuple_GET_SIZE(pool);
2376 r = n;
2377 if (robj != Py_None) {
2378 if (!PyLong_Check(robj)) {
2379 PyErr_SetString(PyExc_TypeError, "Expected int as r");
2380 goto error;
2382 r = PyLong_AsSsize_t(robj);
2383 if (r == -1 && PyErr_Occurred())
2384 goto error;
2386 if (r < 0) {
2387 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2388 goto error;
2391 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2392 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2393 if (indices == NULL || cycles == NULL) {
2394 PyErr_NoMemory();
2395 goto error;
2398 for (i=0 ; i<n ; i++)
2399 indices[i] = i;
2400 for (i=0 ; i<r ; i++)
2401 cycles[i] = n - i;
2403 /* create permutationsobject structure */
2404 po = (permutationsobject *)type->tp_alloc(type, 0);
2405 if (po == NULL)
2406 goto error;
2408 po->pool = pool;
2409 po->indices = indices;
2410 po->cycles = cycles;
2411 po->result = NULL;
2412 po->r = r;
2413 po->stopped = r > n ? 1 : 0;
2415 return (PyObject *)po;
2417 error:
2418 if (indices != NULL)
2419 PyMem_Free(indices);
2420 if (cycles != NULL)
2421 PyMem_Free(cycles);
2422 Py_XDECREF(pool);
2423 return NULL;
2426 static void
2427 permutations_dealloc(permutationsobject *po)
2429 PyObject_GC_UnTrack(po);
2430 Py_XDECREF(po->pool);
2431 Py_XDECREF(po->result);
2432 PyMem_Free(po->indices);
2433 PyMem_Free(po->cycles);
2434 Py_TYPE(po)->tp_free(po);
2437 static int
2438 permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2440 Py_VISIT(po->pool);
2441 Py_VISIT(po->result);
2442 return 0;
2445 static PyObject *
2446 permutations_next(permutationsobject *po)
2448 PyObject *elem;
2449 PyObject *oldelem;
2450 PyObject *pool = po->pool;
2451 Py_ssize_t *indices = po->indices;
2452 Py_ssize_t *cycles = po->cycles;
2453 PyObject *result = po->result;
2454 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2455 Py_ssize_t r = po->r;
2456 Py_ssize_t i, j, k, index;
2458 if (po->stopped)
2459 return NULL;
2461 if (result == NULL) {
2462 /* On the first pass, initialize result tuple using the indices */
2463 result = PyTuple_New(r);
2464 if (result == NULL)
2465 goto empty;
2466 po->result = result;
2467 for (i=0; i<r ; i++) {
2468 index = indices[i];
2469 elem = PyTuple_GET_ITEM(pool, index);
2470 Py_INCREF(elem);
2471 PyTuple_SET_ITEM(result, i, elem);
2473 } else {
2474 if (n == 0)
2475 goto empty;
2477 /* Copy the previous result tuple or re-use it if available */
2478 if (Py_REFCNT(result) > 1) {
2479 PyObject *old_result = result;
2480 result = PyTuple_New(r);
2481 if (result == NULL)
2482 goto empty;
2483 po->result = result;
2484 for (i=0; i<r ; i++) {
2485 elem = PyTuple_GET_ITEM(old_result, i);
2486 Py_INCREF(elem);
2487 PyTuple_SET_ITEM(result, i, elem);
2489 Py_DECREF(old_result);
2491 /* Now, we've got the only copy so we can update it in-place */
2492 assert(r == 0 || Py_REFCNT(result) == 1);
2494 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2495 for (i=r-1 ; i>=0 ; i--) {
2496 cycles[i] -= 1;
2497 if (cycles[i] == 0) {
2498 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2499 index = indices[i];
2500 for (j=i ; j<n-1 ; j++)
2501 indices[j] = indices[j+1];
2502 indices[n-1] = index;
2503 cycles[i] = n - i;
2504 } else {
2505 j = cycles[i];
2506 index = indices[i];
2507 indices[i] = indices[n-j];
2508 indices[n-j] = index;
2510 for (k=i; k<r ; k++) {
2511 /* start with i, the leftmost element that changed */
2512 /* yield tuple(pool[k] for k in indices[:r]) */
2513 index = indices[k];
2514 elem = PyTuple_GET_ITEM(pool, index);
2515 Py_INCREF(elem);
2516 oldelem = PyTuple_GET_ITEM(result, k);
2517 PyTuple_SET_ITEM(result, k, elem);
2518 Py_DECREF(oldelem);
2520 break;
2523 /* If i is negative, then the cycles have all
2524 rolled-over and we're done. */
2525 if (i < 0)
2526 goto empty;
2528 Py_INCREF(result);
2529 return result;
2531 empty:
2532 po->stopped = 1;
2533 return NULL;
2536 PyDoc_STRVAR(permutations_doc,
2537 "permutations(iterable[, r]) --> permutations object\n\
2539 Return successive r-length permutations of elements in the iterable.\n\n\
2540 permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
2542 static PyTypeObject permutations_type = {
2543 PyVarObject_HEAD_INIT(NULL, 0)
2544 "itertools.permutations", /* tp_name */
2545 sizeof(permutationsobject), /* tp_basicsize */
2546 0, /* tp_itemsize */
2547 /* methods */
2548 (destructor)permutations_dealloc, /* tp_dealloc */
2549 0, /* tp_print */
2550 0, /* tp_getattr */
2551 0, /* tp_setattr */
2552 0, /* tp_reserved */
2553 0, /* tp_repr */
2554 0, /* tp_as_number */
2555 0, /* tp_as_sequence */
2556 0, /* tp_as_mapping */
2557 0, /* tp_hash */
2558 0, /* tp_call */
2559 0, /* tp_str */
2560 PyObject_GenericGetAttr, /* tp_getattro */
2561 0, /* tp_setattro */
2562 0, /* tp_as_buffer */
2563 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2564 Py_TPFLAGS_BASETYPE, /* tp_flags */
2565 permutations_doc, /* tp_doc */
2566 (traverseproc)permutations_traverse, /* tp_traverse */
2567 0, /* tp_clear */
2568 0, /* tp_richcompare */
2569 0, /* tp_weaklistoffset */
2570 PyObject_SelfIter, /* tp_iter */
2571 (iternextfunc)permutations_next, /* tp_iternext */
2572 0, /* tp_methods */
2573 0, /* tp_members */
2574 0, /* tp_getset */
2575 0, /* tp_base */
2576 0, /* tp_dict */
2577 0, /* tp_descr_get */
2578 0, /* tp_descr_set */
2579 0, /* tp_dictoffset */
2580 0, /* tp_init */
2581 0, /* tp_alloc */
2582 permutations_new, /* tp_new */
2583 PyObject_GC_Del, /* tp_free */
2587 /* compress object ************************************************************/
2589 /* Equivalent to:
2591 def compress(data, selectors):
2592 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2593 return (d for d, s in zip(data, selectors) if s)
2596 typedef struct {
2597 PyObject_HEAD
2598 PyObject *data;
2599 PyObject *selectors;
2600 } compressobject;
2602 static PyTypeObject compress_type;
2604 static PyObject *
2605 compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2607 PyObject *seq1, *seq2;
2608 PyObject *data=NULL, *selectors=NULL;
2609 compressobject *lz;
2610 static char *kwargs[] = {"data", "selectors", NULL};
2612 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2613 return NULL;
2615 data = PyObject_GetIter(seq1);
2616 if (data == NULL)
2617 goto fail;
2618 selectors = PyObject_GetIter(seq2);
2619 if (selectors == NULL)
2620 goto fail;
2622 /* create compressobject structure */
2623 lz = (compressobject *)type->tp_alloc(type, 0);
2624 if (lz == NULL)
2625 goto fail;
2626 lz->data = data;
2627 lz->selectors = selectors;
2628 return (PyObject *)lz;
2630 fail:
2631 Py_XDECREF(data);
2632 Py_XDECREF(selectors);
2633 return NULL;
2636 static void
2637 compress_dealloc(compressobject *lz)
2639 PyObject_GC_UnTrack(lz);
2640 Py_XDECREF(lz->data);
2641 Py_XDECREF(lz->selectors);
2642 Py_TYPE(lz)->tp_free(lz);
2645 static int
2646 compress_traverse(compressobject *lz, visitproc visit, void *arg)
2648 Py_VISIT(lz->data);
2649 Py_VISIT(lz->selectors);
2650 return 0;
2653 static PyObject *
2654 compress_next(compressobject *lz)
2656 PyObject *data = lz->data, *selectors = lz->selectors;
2657 PyObject *datum, *selector;
2658 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2659 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2660 int ok;
2662 while (1) {
2663 /* Steps: get datum, get selector, evaluate selector.
2664 Order is important (to match the pure python version
2665 in terms of which input gets a chance to raise an
2666 exception first).
2669 datum = datanext(data);
2670 if (datum == NULL)
2671 return NULL;
2673 selector = selectornext(selectors);
2674 if (selector == NULL) {
2675 Py_DECREF(datum);
2676 return NULL;
2679 ok = PyObject_IsTrue(selector);
2680 Py_DECREF(selector);
2681 if (ok == 1)
2682 return datum;
2683 Py_DECREF(datum);
2684 if (ok == -1)
2685 return NULL;
2689 PyDoc_STRVAR(compress_doc,
2690 "compress(data, selectors) --> iterator over selected data\n\
2692 Return data elements corresponding to true selector elements.\n\
2693 Forms a shorter iterator from selected data elements using the\n\
2694 selectors to choose the data elements.");
2696 static PyTypeObject compress_type = {
2697 PyVarObject_HEAD_INIT(NULL, 0)
2698 "itertools.compress", /* tp_name */
2699 sizeof(compressobject), /* tp_basicsize */
2700 0, /* tp_itemsize */
2701 /* methods */
2702 (destructor)compress_dealloc, /* tp_dealloc */
2703 0, /* tp_print */
2704 0, /* tp_getattr */
2705 0, /* tp_setattr */
2706 0, /* tp_reserved */
2707 0, /* tp_repr */
2708 0, /* tp_as_number */
2709 0, /* tp_as_sequence */
2710 0, /* tp_as_mapping */
2711 0, /* tp_hash */
2712 0, /* tp_call */
2713 0, /* tp_str */
2714 PyObject_GenericGetAttr, /* tp_getattro */
2715 0, /* tp_setattro */
2716 0, /* tp_as_buffer */
2717 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2718 Py_TPFLAGS_BASETYPE, /* tp_flags */
2719 compress_doc, /* tp_doc */
2720 (traverseproc)compress_traverse, /* tp_traverse */
2721 0, /* tp_clear */
2722 0, /* tp_richcompare */
2723 0, /* tp_weaklistoffset */
2724 PyObject_SelfIter, /* tp_iter */
2725 (iternextfunc)compress_next, /* tp_iternext */
2726 0, /* tp_methods */
2727 0, /* tp_members */
2728 0, /* tp_getset */
2729 0, /* tp_base */
2730 0, /* tp_dict */
2731 0, /* tp_descr_get */
2732 0, /* tp_descr_set */
2733 0, /* tp_dictoffset */
2734 0, /* tp_init */
2735 0, /* tp_alloc */
2736 compress_new, /* tp_new */
2737 PyObject_GC_Del, /* tp_free */
2741 /* filterfalse object ************************************************************/
2743 typedef struct {
2744 PyObject_HEAD
2745 PyObject *func;
2746 PyObject *it;
2747 } filterfalseobject;
2749 static PyTypeObject filterfalse_type;
2751 static PyObject *
2752 filterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2754 PyObject *func, *seq;
2755 PyObject *it;
2756 filterfalseobject *lz;
2758 if (type == &filterfalse_type &&
2759 !_PyArg_NoKeywords("filterfalse()", kwds))
2760 return NULL;
2762 if (!PyArg_UnpackTuple(args, "filterfalse", 2, 2, &func, &seq))
2763 return NULL;
2765 /* Get iterator. */
2766 it = PyObject_GetIter(seq);
2767 if (it == NULL)
2768 return NULL;
2770 /* create filterfalseobject structure */
2771 lz = (filterfalseobject *)type->tp_alloc(type, 0);
2772 if (lz == NULL) {
2773 Py_DECREF(it);
2774 return NULL;
2776 Py_INCREF(func);
2777 lz->func = func;
2778 lz->it = it;
2780 return (PyObject *)lz;
2783 static void
2784 filterfalse_dealloc(filterfalseobject *lz)
2786 PyObject_GC_UnTrack(lz);
2787 Py_XDECREF(lz->func);
2788 Py_XDECREF(lz->it);
2789 Py_TYPE(lz)->tp_free(lz);
2792 static int
2793 filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
2795 Py_VISIT(lz->it);
2796 Py_VISIT(lz->func);
2797 return 0;
2800 static PyObject *
2801 filterfalse_next(filterfalseobject *lz)
2803 PyObject *item;
2804 PyObject *it = lz->it;
2805 long ok;
2806 PyObject *(*iternext)(PyObject *);
2808 iternext = *Py_TYPE(it)->tp_iternext;
2809 for (;;) {
2810 item = iternext(it);
2811 if (item == NULL)
2812 return NULL;
2814 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2815 ok = PyObject_IsTrue(item);
2816 } else {
2817 PyObject *good;
2818 good = PyObject_CallFunctionObjArgs(lz->func,
2819 item, NULL);
2820 if (good == NULL) {
2821 Py_DECREF(item);
2822 return NULL;
2824 ok = PyObject_IsTrue(good);
2825 Py_DECREF(good);
2827 if (!ok)
2828 return item;
2829 Py_DECREF(item);
2833 PyDoc_STRVAR(filterfalse_doc,
2834 "filterfalse(function or None, sequence) --> filterfalse object\n\
2836 Return those items of sequence for which function(item) is false.\n\
2837 If function is None, return the items that are false.");
2839 static PyTypeObject filterfalse_type = {
2840 PyVarObject_HEAD_INIT(NULL, 0)
2841 "itertools.filterfalse", /* tp_name */
2842 sizeof(filterfalseobject), /* tp_basicsize */
2843 0, /* tp_itemsize */
2844 /* methods */
2845 (destructor)filterfalse_dealloc, /* tp_dealloc */
2846 0, /* tp_print */
2847 0, /* tp_getattr */
2848 0, /* tp_setattr */
2849 0, /* tp_reserved */
2850 0, /* tp_repr */
2851 0, /* tp_as_number */
2852 0, /* tp_as_sequence */
2853 0, /* tp_as_mapping */
2854 0, /* tp_hash */
2855 0, /* tp_call */
2856 0, /* tp_str */
2857 PyObject_GenericGetAttr, /* tp_getattro */
2858 0, /* tp_setattro */
2859 0, /* tp_as_buffer */
2860 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2861 Py_TPFLAGS_BASETYPE, /* tp_flags */
2862 filterfalse_doc, /* tp_doc */
2863 (traverseproc)filterfalse_traverse, /* tp_traverse */
2864 0, /* tp_clear */
2865 0, /* tp_richcompare */
2866 0, /* tp_weaklistoffset */
2867 PyObject_SelfIter, /* tp_iter */
2868 (iternextfunc)filterfalse_next, /* tp_iternext */
2869 0, /* tp_methods */
2870 0, /* tp_members */
2871 0, /* tp_getset */
2872 0, /* tp_base */
2873 0, /* tp_dict */
2874 0, /* tp_descr_get */
2875 0, /* tp_descr_set */
2876 0, /* tp_dictoffset */
2877 0, /* tp_init */
2878 0, /* tp_alloc */
2879 filterfalse_new, /* tp_new */
2880 PyObject_GC_Del, /* tp_free */
2884 /* count object ************************************************************/
2886 typedef struct {
2887 PyObject_HEAD
2888 Py_ssize_t cnt;
2889 PyObject *long_cnt;
2890 PyObject *long_step;
2891 } countobject;
2893 /* Counting logic and invariants:
2895 fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
2897 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
2898 Advances with: cnt += 1
2899 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
2901 slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
2903 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
2904 All counting is done with python objects (no overflows or underflows).
2905 Advances with: long_cnt += long_step
2906 Step may be zero -- effectively a slow version of repeat(cnt).
2907 Either long_cnt or long_step may be a float, Fraction, or Decimal.
2910 static PyTypeObject count_type;
2912 static PyObject *
2913 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2915 countobject *lz;
2916 int slow_mode = 0;
2917 Py_ssize_t cnt = 0;
2918 PyObject *long_cnt = NULL;
2919 PyObject *long_step = NULL;
2920 static char *kwlist[] = {"start", "step", 0};
2922 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
2923 kwlist, &long_cnt, &long_step))
2924 return NULL;
2926 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
2927 (long_step != NULL && !PyNumber_Check(long_step))) {
2928 PyErr_SetString(PyExc_TypeError, "a number is required");
2929 return NULL;
2932 if (long_cnt != NULL) {
2933 cnt = PyLong_AsSsize_t(long_cnt);
2934 if ((cnt == -1 && PyErr_Occurred()) || !PyLong_Check(long_cnt)) {
2935 PyErr_Clear();
2936 slow_mode = 1;
2938 Py_INCREF(long_cnt);
2939 } else {
2940 cnt = 0;
2941 long_cnt = PyLong_FromLong(0);
2944 /* If not specified, step defaults to 1 */
2945 if (long_step == NULL) {
2946 long_step = PyLong_FromLong(1);
2947 if (long_step == NULL) {
2948 Py_DECREF(long_cnt);
2949 return NULL;
2951 } else
2952 Py_INCREF(long_step);
2954 assert(long_cnt != NULL && long_step != NULL);
2956 /* Fast mode only works when the step is 1 */
2957 if (!PyLong_Check(long_step) ||
2958 PyLong_AS_LONG(long_step) != 1) {
2959 slow_mode = 1;
2962 if (slow_mode)
2963 cnt = PY_SSIZE_T_MAX;
2964 else
2965 Py_CLEAR(long_cnt);
2967 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
2968 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
2969 assert(slow_mode ||
2970 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
2972 /* create countobject structure */
2973 lz = (countobject *)type->tp_alloc(type, 0);
2974 if (lz == NULL) {
2975 Py_XDECREF(long_cnt);
2976 return NULL;
2978 lz->cnt = cnt;
2979 lz->long_cnt = long_cnt;
2980 lz->long_step = long_step;
2982 return (PyObject *)lz;
2985 static void
2986 count_dealloc(countobject *lz)
2988 PyObject_GC_UnTrack(lz);
2989 Py_XDECREF(lz->long_cnt);
2990 Py_XDECREF(lz->long_step);
2991 Py_TYPE(lz)->tp_free(lz);
2994 static int
2995 count_traverse(countobject *lz, visitproc visit, void *arg)
2997 Py_VISIT(lz->long_cnt);
2998 Py_VISIT(lz->long_step);
2999 return 0;
3002 static PyObject *
3003 count_nextlong(countobject *lz)
3005 PyObject *long_cnt;
3006 PyObject *stepped_up;
3008 long_cnt = lz->long_cnt;
3009 if (long_cnt == NULL) {
3010 /* Switch to slow_mode */
3011 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
3012 if (long_cnt == NULL)
3013 return NULL;
3015 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
3017 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3018 if (stepped_up == NULL)
3019 return NULL;
3020 lz->long_cnt = stepped_up;
3021 return long_cnt;
3024 static PyObject *
3025 count_next(countobject *lz)
3027 if (lz->cnt == PY_SSIZE_T_MAX)
3028 return count_nextlong(lz);
3029 return PyLong_FromSsize_t(lz->cnt++);
3032 static PyObject *
3033 count_repr(countobject *lz)
3035 if (lz->cnt != PY_SSIZE_T_MAX)
3036 return PyUnicode_FromFormat("count(%zd)", lz->cnt);
3038 if (PyLong_Check(lz->long_step)) {
3039 long step = PyLong_AsLong(lz->long_step);
3040 if (step == -1 && PyErr_Occurred()) {
3041 PyErr_Clear();
3043 if (step == 1) {
3044 /* Don't display step when it is an integer equal to 1 */
3045 return PyUnicode_FromFormat("count(%R)", lz->long_cnt);
3048 return PyUnicode_FromFormat("count(%R, %R)",
3049 lz->long_cnt, lz->long_step);
3052 static PyObject *
3053 count_reduce(countobject *lz)
3055 if (lz->cnt == PY_SSIZE_T_MAX)
3056 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
3057 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
3060 PyDoc_STRVAR(count_reduce_doc, "Return state information for pickling.");
3062 static PyMethodDef count_methods[] = {
3063 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
3064 count_reduce_doc},
3065 {NULL, NULL} /* sentinel */
3068 PyDoc_STRVAR(count_doc,
3069 "count(start=0, step=1) --> count object\n\
3071 Return a count object whose .__next__() method returns consecutive values.\n\
3072 Equivalent to:\n\n\
3073 def count(firstval=0, step=1):\n\
3074 x = firstval\n\
3075 while 1:\n\
3076 yield x\n\
3077 x += step\n");
3079 static PyTypeObject count_type = {
3080 PyVarObject_HEAD_INIT(NULL, 0)
3081 "itertools.count", /* tp_name */
3082 sizeof(countobject), /* tp_basicsize */
3083 0, /* tp_itemsize */
3084 /* methods */
3085 (destructor)count_dealloc, /* tp_dealloc */
3086 0, /* tp_print */
3087 0, /* tp_getattr */
3088 0, /* tp_setattr */
3089 0, /* tp_reserved */
3090 (reprfunc)count_repr, /* tp_repr */
3091 0, /* tp_as_number */
3092 0, /* tp_as_sequence */
3093 0, /* tp_as_mapping */
3094 0, /* tp_hash */
3095 0, /* tp_call */
3096 0, /* tp_str */
3097 PyObject_GenericGetAttr, /* tp_getattro */
3098 0, /* tp_setattro */
3099 0, /* tp_as_buffer */
3100 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3101 Py_TPFLAGS_BASETYPE, /* tp_flags */
3102 count_doc, /* tp_doc */
3103 (traverseproc)count_traverse, /* tp_traverse */
3104 0, /* tp_clear */
3105 0, /* tp_richcompare */
3106 0, /* tp_weaklistoffset */
3107 PyObject_SelfIter, /* tp_iter */
3108 (iternextfunc)count_next, /* tp_iternext */
3109 count_methods, /* tp_methods */
3110 0, /* tp_members */
3111 0, /* tp_getset */
3112 0, /* tp_base */
3113 0, /* tp_dict */
3114 0, /* tp_descr_get */
3115 0, /* tp_descr_set */
3116 0, /* tp_dictoffset */
3117 0, /* tp_init */
3118 0, /* tp_alloc */
3119 count_new, /* tp_new */
3120 PyObject_GC_Del, /* tp_free */
3124 /* repeat object ************************************************************/
3126 typedef struct {
3127 PyObject_HEAD
3128 PyObject *element;
3129 Py_ssize_t cnt;
3130 } repeatobject;
3132 static PyTypeObject repeat_type;
3134 static PyObject *
3135 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3137 repeatobject *ro;
3138 PyObject *element;
3139 Py_ssize_t cnt = -1;
3140 static char *kwargs[] = {"object", "times", NULL};
3142 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3143 &element, &cnt))
3144 return NULL;
3146 if (PyTuple_Size(args) == 2 && cnt < 0)
3147 cnt = 0;
3149 ro = (repeatobject *)type->tp_alloc(type, 0);
3150 if (ro == NULL)
3151 return NULL;
3152 Py_INCREF(element);
3153 ro->element = element;
3154 ro->cnt = cnt;
3155 return (PyObject *)ro;
3158 static void
3159 repeat_dealloc(repeatobject *ro)
3161 PyObject_GC_UnTrack(ro);
3162 Py_XDECREF(ro->element);
3163 Py_TYPE(ro)->tp_free(ro);
3166 static int
3167 repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3169 Py_VISIT(ro->element);
3170 return 0;
3173 static PyObject *
3174 repeat_next(repeatobject *ro)
3176 if (ro->cnt == 0)
3177 return NULL;
3178 if (ro->cnt > 0)
3179 ro->cnt--;
3180 Py_INCREF(ro->element);
3181 return ro->element;
3184 static PyObject *
3185 repeat_repr(repeatobject *ro)
3187 if (ro->cnt == -1)
3188 return PyUnicode_FromFormat("repeat(%R)", ro->element);
3189 else
3190 return PyUnicode_FromFormat("repeat(%R, %zd)", ro->element, ro->cnt);
3193 static PyObject *
3194 repeat_len(repeatobject *ro)
3196 if (ro->cnt == -1) {
3197 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3198 return NULL;
3200 return PyLong_FromSize_t(ro->cnt);
3203 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
3205 static PyMethodDef repeat_methods[] = {
3206 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3207 {NULL, NULL} /* sentinel */
3210 PyDoc_STRVAR(repeat_doc,
3211 "repeat(object [,times]) -> create an iterator which returns the object\n\
3212 for the specified number of times. If not specified, returns the object\n\
3213 endlessly.");
3215 static PyTypeObject repeat_type = {
3216 PyVarObject_HEAD_INIT(NULL, 0)
3217 "itertools.repeat", /* tp_name */
3218 sizeof(repeatobject), /* tp_basicsize */
3219 0, /* tp_itemsize */
3220 /* methods */
3221 (destructor)repeat_dealloc, /* tp_dealloc */
3222 0, /* tp_print */
3223 0, /* tp_getattr */
3224 0, /* tp_setattr */
3225 0, /* tp_reserved */
3226 (reprfunc)repeat_repr, /* tp_repr */
3227 0, /* tp_as_number */
3228 0, /* tp_as_sequence */
3229 0, /* tp_as_mapping */
3230 0, /* tp_hash */
3231 0, /* tp_call */
3232 0, /* tp_str */
3233 PyObject_GenericGetAttr, /* tp_getattro */
3234 0, /* tp_setattro */
3235 0, /* tp_as_buffer */
3236 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3237 Py_TPFLAGS_BASETYPE, /* tp_flags */
3238 repeat_doc, /* tp_doc */
3239 (traverseproc)repeat_traverse, /* tp_traverse */
3240 0, /* tp_clear */
3241 0, /* tp_richcompare */
3242 0, /* tp_weaklistoffset */
3243 PyObject_SelfIter, /* tp_iter */
3244 (iternextfunc)repeat_next, /* tp_iternext */
3245 repeat_methods, /* tp_methods */
3246 0, /* tp_members */
3247 0, /* tp_getset */
3248 0, /* tp_base */
3249 0, /* tp_dict */
3250 0, /* tp_descr_get */
3251 0, /* tp_descr_set */
3252 0, /* tp_dictoffset */
3253 0, /* tp_init */
3254 0, /* tp_alloc */
3255 repeat_new, /* tp_new */
3256 PyObject_GC_Del, /* tp_free */
3259 /* ziplongest object ************************************************************/
3261 #include "Python.h"
3263 typedef struct {
3264 PyObject_HEAD
3265 Py_ssize_t tuplesize;
3266 Py_ssize_t numactive;
3267 PyObject *ittuple; /* tuple of iterators */
3268 PyObject *result;
3269 PyObject *fillvalue;
3270 } ziplongestobject;
3272 static PyTypeObject ziplongest_type;
3274 static PyObject *
3275 zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3277 ziplongestobject *lz;
3278 Py_ssize_t i;
3279 PyObject *ittuple; /* tuple of iterators */
3280 PyObject *result;
3281 PyObject *fillvalue = Py_None;
3282 Py_ssize_t tuplesize = PySequence_Length(args);
3284 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3285 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3286 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3287 PyErr_SetString(PyExc_TypeError,
3288 "zip_longest() got an unexpected keyword argument");
3289 return NULL;
3293 /* args must be a tuple */
3294 assert(PyTuple_Check(args));
3296 /* obtain iterators */
3297 ittuple = PyTuple_New(tuplesize);
3298 if (ittuple == NULL)
3299 return NULL;
3300 for (i=0; i < tuplesize; ++i) {
3301 PyObject *item = PyTuple_GET_ITEM(args, i);
3302 PyObject *it = PyObject_GetIter(item);
3303 if (it == NULL) {
3304 if (PyErr_ExceptionMatches(PyExc_TypeError))
3305 PyErr_Format(PyExc_TypeError,
3306 "zip_longest argument #%zd must support iteration",
3307 i+1);
3308 Py_DECREF(ittuple);
3309 return NULL;
3311 PyTuple_SET_ITEM(ittuple, i, it);
3314 /* create a result holder */
3315 result = PyTuple_New(tuplesize);
3316 if (result == NULL) {
3317 Py_DECREF(ittuple);
3318 return NULL;
3320 for (i=0 ; i < tuplesize ; i++) {
3321 Py_INCREF(Py_None);
3322 PyTuple_SET_ITEM(result, i, Py_None);
3325 /* create ziplongestobject structure */
3326 lz = (ziplongestobject *)type->tp_alloc(type, 0);
3327 if (lz == NULL) {
3328 Py_DECREF(ittuple);
3329 Py_DECREF(result);
3330 return NULL;
3332 lz->ittuple = ittuple;
3333 lz->tuplesize = tuplesize;
3334 lz->numactive = tuplesize;
3335 lz->result = result;
3336 Py_INCREF(fillvalue);
3337 lz->fillvalue = fillvalue;
3338 return (PyObject *)lz;
3341 static void
3342 zip_longest_dealloc(ziplongestobject *lz)
3344 PyObject_GC_UnTrack(lz);
3345 Py_XDECREF(lz->ittuple);
3346 Py_XDECREF(lz->result);
3347 Py_XDECREF(lz->fillvalue);
3348 Py_TYPE(lz)->tp_free(lz);
3351 static int
3352 zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
3354 Py_VISIT(lz->ittuple);
3355 Py_VISIT(lz->result);
3356 Py_VISIT(lz->fillvalue);
3357 return 0;
3360 static PyObject *
3361 zip_longest_next(ziplongestobject *lz)
3363 Py_ssize_t i;
3364 Py_ssize_t tuplesize = lz->tuplesize;
3365 PyObject *result = lz->result;
3366 PyObject *it;
3367 PyObject *item;
3368 PyObject *olditem;
3370 if (tuplesize == 0)
3371 return NULL;
3372 if (lz->numactive == 0)
3373 return NULL;
3374 if (Py_REFCNT(result) == 1) {
3375 Py_INCREF(result);
3376 for (i=0 ; i < tuplesize ; i++) {
3377 it = PyTuple_GET_ITEM(lz->ittuple, i);
3378 if (it == NULL) {
3379 Py_INCREF(lz->fillvalue);
3380 item = lz->fillvalue;
3381 } else {
3382 item = PyIter_Next(it);
3383 if (item == NULL) {
3384 lz->numactive -= 1;
3385 if (lz->numactive == 0 || PyErr_Occurred()) {
3386 lz->numactive = 0;
3387 Py_DECREF(result);
3388 return NULL;
3389 } else {
3390 Py_INCREF(lz->fillvalue);
3391 item = lz->fillvalue;
3392 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3393 Py_DECREF(it);
3397 olditem = PyTuple_GET_ITEM(result, i);
3398 PyTuple_SET_ITEM(result, i, item);
3399 Py_DECREF(olditem);
3401 } else {
3402 result = PyTuple_New(tuplesize);
3403 if (result == NULL)
3404 return NULL;
3405 for (i=0 ; i < tuplesize ; i++) {
3406 it = PyTuple_GET_ITEM(lz->ittuple, i);
3407 if (it == NULL) {
3408 Py_INCREF(lz->fillvalue);
3409 item = lz->fillvalue;
3410 } else {
3411 item = PyIter_Next(it);
3412 if (item == NULL) {
3413 lz->numactive -= 1;
3414 if (lz->numactive == 0 || PyErr_Occurred()) {
3415 lz->numactive = 0;
3416 Py_DECREF(result);
3417 return NULL;
3418 } else {
3419 Py_INCREF(lz->fillvalue);
3420 item = lz->fillvalue;
3421 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3422 Py_DECREF(it);
3426 PyTuple_SET_ITEM(result, i, item);
3429 return result;
3432 PyDoc_STRVAR(zip_longest_doc,
3433 "zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
3435 Return an zip_longest object whose .__next__() method returns a tuple where\n\
3436 the i-th element comes from the i-th iterable argument. The .__next__()\n\
3437 method continues until the longest iterable in the argument sequence\n\
3438 is exhausted and then it raises StopIteration. When the shorter iterables\n\
3439 are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3440 defaults to None or can be specified by a keyword argument.\n\
3443 static PyTypeObject ziplongest_type = {
3444 PyVarObject_HEAD_INIT(NULL, 0)
3445 "itertools.zip_longest", /* tp_name */
3446 sizeof(ziplongestobject), /* tp_basicsize */
3447 0, /* tp_itemsize */
3448 /* methods */
3449 (destructor)zip_longest_dealloc, /* tp_dealloc */
3450 0, /* tp_print */
3451 0, /* tp_getattr */
3452 0, /* tp_setattr */
3453 0, /* tp_reserved */
3454 0, /* tp_repr */
3455 0, /* tp_as_number */
3456 0, /* tp_as_sequence */
3457 0, /* tp_as_mapping */
3458 0, /* tp_hash */
3459 0, /* tp_call */
3460 0, /* tp_str */
3461 PyObject_GenericGetAttr, /* tp_getattro */
3462 0, /* tp_setattro */
3463 0, /* tp_as_buffer */
3464 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3465 Py_TPFLAGS_BASETYPE, /* tp_flags */
3466 zip_longest_doc, /* tp_doc */
3467 (traverseproc)zip_longest_traverse, /* tp_traverse */
3468 0, /* tp_clear */
3469 0, /* tp_richcompare */
3470 0, /* tp_weaklistoffset */
3471 PyObject_SelfIter, /* tp_iter */
3472 (iternextfunc)zip_longest_next, /* tp_iternext */
3473 0, /* tp_methods */
3474 0, /* tp_members */
3475 0, /* tp_getset */
3476 0, /* tp_base */
3477 0, /* tp_dict */
3478 0, /* tp_descr_get */
3479 0, /* tp_descr_set */
3480 0, /* tp_dictoffset */
3481 0, /* tp_init */
3482 0, /* tp_alloc */
3483 zip_longest_new, /* tp_new */
3484 PyObject_GC_Del, /* tp_free */
3487 /* module level code ********************************************************/
3489 PyDoc_STRVAR(module_doc,
3490 "Functional tools for creating and using iterators.\n\
3492 Infinite iterators:\n\
3493 count([n]) --> n, n+1, n+2, ...\n\
3494 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
3495 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
3497 Iterators terminating on the shortest input sequence:\n\
3498 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3499 compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
3500 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3501 groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
3502 filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
3503 islice(seq, [start,] stop [, step]) --> elements from\n\
3504 seq[start:stop:step]\n\
3505 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
3506 tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
3507 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3508 zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3510 Combinatoric generators:\n\
3511 product(p, q, ... [repeat=1]) --> cartesian product\n\
3512 permutations(p[, r])\n\
3513 combinations(p, r)\n\
3514 combinations_with_replacement(p, r)\n\
3518 static PyMethodDef module_methods[] = {
3519 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3520 {NULL, NULL} /* sentinel */
3524 static struct PyModuleDef itertoolsmodule = {
3525 PyModuleDef_HEAD_INIT,
3526 "itertools",
3527 module_doc,
3529 module_methods,
3530 NULL,
3531 NULL,
3532 NULL,
3533 NULL
3536 PyMODINIT_FUNC
3537 PyInit_itertools(void)
3539 int i;
3540 PyObject *m;
3541 char *name;
3542 PyTypeObject *typelist[] = {
3543 &combinations_type,
3544 &cwr_type,
3545 &cycle_type,
3546 &dropwhile_type,
3547 &takewhile_type,
3548 &islice_type,
3549 &starmap_type,
3550 &chain_type,
3551 &compress_type,
3552 &filterfalse_type,
3553 &count_type,
3554 &ziplongest_type,
3555 &permutations_type,
3556 &product_type,
3557 &repeat_type,
3558 &groupby_type,
3559 NULL
3562 Py_TYPE(&teedataobject_type) = &PyType_Type;
3563 m = PyModule_Create(&itertoolsmodule);
3564 if (m == NULL)
3565 return NULL;
3567 for (i=0 ; typelist[i] != NULL ; i++) {
3568 if (PyType_Ready(typelist[i]) < 0)
3569 return NULL;
3570 name = strchr(typelist[i]->tp_name, '.');
3571 assert (name != NULL);
3572 Py_INCREF(typelist[i]);
3573 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3576 if (PyType_Ready(&teedataobject_type) < 0)
3577 return NULL;
3578 if (PyType_Ready(&tee_type) < 0)
3579 return NULL;
3580 if (PyType_Ready(&_grouper_type) < 0)
3581 return NULL;
3582 return m;