?  search:
summary | log | graphiclog1 | graphiclog2 | commit | commitdiff | tree | refs | edit | fork
blame | history | raw | HEAD
Fix a refleak introduced by r66677.
[python.git] / Modules / itertoolsmodule.c
blob18a1229b14f730b6541730f815acf73a5efaba9a
1
2 #include "Python.h"
3 #include "structmember.h"
4
5 /* Itertools module written and maintained
6 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9 */
10
11
12 /* groupby object ***********************************************************/
13
14 typedef struct {
15 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
21 } groupbyobject;
22
23 static PyTypeObject groupby_type;
24 static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26 static PyObject *
27 groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28 {
29 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
32
33 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
51 }
52
53 static void
54 groupby_dealloc(groupbyobject *gbo)
55 {
56 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 Py_TYPE(gbo)->tp_free(gbo);
63 }
64
65 static int
66 groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67 {
68 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
73 return 0;
74 }
75
76 static PyObject *
77 groupby_next(groupbyobject *gbo)
78 {
79 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
80
81 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
89
90 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
97
98 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
101
102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
113
114 tmp = gbo->currkey;
115 gbo->currkey = newkey;
116 Py_XDECREF(tmp);
117
118 tmp = gbo->currvalue;
119 gbo->currvalue = newvalue;
120 Py_XDECREF(tmp);
121 }
122
123 Py_INCREF(gbo->currkey);
124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
127
128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
131
132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
135 }
136
137 PyDoc_STRVAR(groupby_doc,
138 "groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139 (key, sub-iterator) grouped by each value of key(value).\n");
140
141 static PyTypeObject groupby_type = {
142 PyVarObject_HEAD_INIT(NULL, 0)
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_compare */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
183 };
184
185
186 /* _grouper object (internal) ************************************************/
187
188 typedef struct {
189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
192 } _grouperobject;
193
194 static PyTypeObject _grouper_type;
195
196 static PyObject *
197 _grouper_create(groupbyobject *parent, PyObject *tgtkey)
198 {
199 _grouperobject *igo;
200
201 igo = PyObject_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);
208
209 PyObject_GC_Track(igo);
210 return (PyObject *)igo;
211 }
212
213 static void
214 _grouper_dealloc(_grouperobject *igo)
215 {
216 PyObject_GC_UnTrack(igo);
217 Py_DECREF(igo->parent);
218 Py_DECREF(igo->tgtkey);
219 PyObject_GC_Del(igo);
220 }
221
222 static int
223 _grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
224 {
225 Py_VISIT(igo->parent);
226 Py_VISIT(igo->tgtkey);
227 return 0;
228 }
229
230 static PyObject *
231 _grouper_next(_grouperobject *igo)
232 {
233 groupbyobject *gbo = (groupbyobject *)igo->parent;
234 PyObject *newvalue, *newkey, *r;
235 int rcmp;
236
237 if (gbo->currvalue == NULL) {
238 newvalue = PyIter_Next(gbo->it);
239 if (newvalue == NULL)
240 return NULL;
241
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;
251 }
252 }
253
254 assert(gbo->currkey == NULL);
255 gbo->currkey = newkey;
256 gbo->currvalue = newvalue;
257 }
258
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;
264
265 r = gbo->currvalue;
266 gbo->currvalue = NULL;
267 Py_CLEAR(gbo->currkey);
268
269 return r;
270 }
271
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_compare */
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 */
313 };
314
315
316
317 /* tee object and with supporting function and objects ***************/
318
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.
326 */
327 #define LINKCELLS 57
328
329 typedef struct {
330 PyObject_HEAD
331 PyObject *it;
332 int numread;
333 PyObject *nextlink;
334 PyObject *(values[LINKCELLS]);
335 } teedataobject;
336
337 typedef struct {
338 PyObject_HEAD
339 teedataobject *dataobj;
340 int index;
341 PyObject *weakreflist;
342 } teeobject;
343
344 static PyTypeObject teedataobject_type;
345
346 static PyObject *
347 teedataobject_new(PyObject *it)
348 {
349 teedataobject *tdo;
350
351 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
352 if (tdo == NULL)
353 return NULL;
354
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;
361 }
362
363 static PyObject *
364 teedataobject_jumplink(teedataobject *tdo)
365 {
366 if (tdo->nextlink == NULL)
367 tdo->nextlink = teedataobject_new(tdo->it);
368 Py_XINCREF(tdo->nextlink);
369 return tdo->nextlink;
370 }
371
372 static PyObject *
373 teedataobject_getitem(teedataobject *tdo, int i)
374 {
375 PyObject *value;
376
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;
388 }
389 Py_INCREF(value);
390 return value;
391 }
392
393 static int
394 teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
395 {
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;
402 }
403
404 static int
405 teedataobject_clear(teedataobject *tdo)
406 {
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;
413 }
414
415 static void
416 teedataobject_dealloc(teedataobject *tdo)
417 {
418 PyObject_GC_UnTrack(tdo);
419 teedataobject_clear(tdo);
420 PyObject_GC_Del(tdo);
421 }
422
423 PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
424
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_compare */
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 */
466 };
467
468
469 static PyTypeObject tee_type;
470
471 static PyObject *
472 tee_next(teeobject *to)
473 {
474 PyObject *value, *link;
475
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;
481 }
482 value = teedataobject_getitem(to->dataobj, to->index);
483 if (value == NULL)
484 return NULL;
485 to->index++;
486 return value;
487 }
488
489 static int
490 tee_traverse(teeobject *to, visitproc visit, void *arg)
491 {
492 Py_VISIT((PyObject *)to->dataobj);
493 return 0;
494 }
495
496 static PyObject *
497 tee_copy(teeobject *to)
498 {
499 teeobject *newto;
500
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;
510 }
511
512 PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
513
514 static PyObject *
515 tee_fromiterable(PyObject *iterable)
516 {
517 teeobject *to;
518 PyObject *it = NULL;
519
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;
526 }
527
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;
536 }
537
538 to->index = 0;
539 to->weakreflist = NULL;
540 PyObject_GC_Track(to);
541 done:
542 Py_XDECREF(it);
543 return (PyObject *)to;
544 }
545
546 static PyObject *
547 tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
548 {
549 PyObject *iterable;
550
551 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
552 return NULL;
553 return tee_fromiterable(iterable);
554 }
555
556 static int
557 tee_clear(teeobject *to)
558 {
559 if (to->weakreflist != NULL)
560 PyObject_ClearWeakRefs((PyObject *) to);
561 Py_CLEAR(to->dataobj);
562 return 0;
563 }
564
565 static void
566 tee_dealloc(teeobject *to)
567 {
568 PyObject_GC_UnTrack(to);
569 tee_clear(to);
570 PyObject_GC_Del(to);
571 }
572
573 PyDoc_STRVAR(teeobject_doc,
574 "Iterator wrapped to make it copyable");
575
576 static PyMethodDef tee_methods[] = {
577 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
578 {NULL, NULL} /* sentinel */
579 };
580
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_compare */
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 */
622 };
623
624 static PyObject *
625 tee(PyObject *self, PyObject *args)
626 {
627 Py_ssize_t i, n=2;
628 PyObject *it, *iterable, *copyable, *result;
629
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;
635 }
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;
645 }
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;
652 }
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;
661 }
662 PyTuple_SET_ITEM(result, i, copyable);
663 }
664 return result;
665 }
666
667 PyDoc_STRVAR(tee_doc,
668 "tee(iterable, n=2) --> tuple of n independent iterators.");
669
670
671 /* cycle object **********************************************************/
672
673 typedef struct {
674 PyObject_HEAD
675 PyObject *it;
676 PyObject *saved;
677 int firstpass;
678 } cycleobject;
679
680 static PyTypeObject cycle_type;
681
682 static PyObject *
683 cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
684 {
685 PyObject *it;
686 PyObject *iterable;
687 PyObject *saved;
688 cycleobject *lz;
689
690 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
691 return NULL;
692
693 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
694 return NULL;
695
696 /* Get iterator. */
697 it = PyObject_GetIter(iterable);
698 if (it == NULL)
699 return NULL;
700
701 saved = PyList_New(0);
702 if (saved == NULL) {
703 Py_DECREF(it);
704 return NULL;
705 }
706
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;
713 }
714 lz->it = it;
715 lz->saved = saved;
716 lz->firstpass = 0;
717
718 return (PyObject *)lz;
719 }
720
721 static void
722 cycle_dealloc(cycleobject *lz)
723 {
724 PyObject_GC_UnTrack(lz);
725 Py_XDECREF(lz->saved);
726 Py_XDECREF(lz->it);
727 Py_TYPE(lz)->tp_free(lz);
728 }
729
730 static int
731 cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
732 {
733 Py_VISIT(lz->it);
734 Py_VISIT(lz->saved);
735 return 0;
736 }
737
738 static PyObject *
739 cycle_next(cycleobject *lz)
740 {
741 PyObject *item;
742 PyObject *it;
743 PyObject *tmp;
744
745 while (1) {
746 item = PyIter_Next(lz->it);
747 if (item != NULL) {
748 if (!lz->firstpass)
749 PyList_Append(lz->saved, item);
750 return item;
751 }
752 if (PyErr_Occurred()) {
753 if (PyErr_ExceptionMatches(PyExc_StopIteration))
754 PyErr_Clear();
755 else
756 return NULL;
757 }
758 if (PyList_Size(lz->saved) == 0)
759 return NULL;
760 it = PyObject_GetIter(lz->saved);
761 if (it == NULL)
762 return NULL;
763 tmp = lz->it;
764 lz->it = it;
765 lz->firstpass = 1;
766 Py_DECREF(tmp);
767 }
768 }
769
770 PyDoc_STRVAR(cycle_doc,
771 "cycle(iterable) --> cycle object\n\
772 \n\
773 Return elements from the iterable until it is exhausted.\n\
774 Then repeat the sequence indefinitely.");
775
776 static PyTypeObject cycle_type = {
777 PyVarObject_HEAD_INIT(NULL, 0)
778 "itertools.cycle", /* tp_name */
779 sizeof(cycleobject), /* tp_basicsize */
780 0, /* tp_itemsize */
781 /* methods */
782 (destructor)cycle_dealloc, /* tp_dealloc */
783 0, /* tp_print */
784 0, /* tp_getattr */
785 0, /* tp_setattr */
786 0, /* tp_compare */
787 0, /* tp_repr */
788 0, /* tp_as_number */
789 0, /* tp_as_sequence */
790 0, /* tp_as_mapping */
791 0, /* tp_hash */
792 0, /* tp_call */
793 0, /* tp_str */
794 PyObject_GenericGetAttr, /* tp_getattro */
795 0, /* tp_setattro */
796 0, /* tp_as_buffer */
797 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
798 Py_TPFLAGS_BASETYPE, /* tp_flags */
799 cycle_doc, /* tp_doc */
800 (traverseproc)cycle_traverse, /* tp_traverse */
801 0, /* tp_clear */
802 0, /* tp_richcompare */
803 0, /* tp_weaklistoffset */
804 PyObject_SelfIter, /* tp_iter */
805 (iternextfunc)cycle_next, /* tp_iternext */
806 0, /* tp_methods */
807 0, /* tp_members */
808 0, /* tp_getset */
809 0, /* tp_base */
810 0, /* tp_dict */
811 0, /* tp_descr_get */
812 0, /* tp_descr_set */
813 0, /* tp_dictoffset */
814 0, /* tp_init */
815 0, /* tp_alloc */
816 cycle_new, /* tp_new */
817 PyObject_GC_Del, /* tp_free */
818 };
819
820
821 /* dropwhile object **********************************************************/
822
823 typedef struct {
824 PyObject_HEAD
825 PyObject *func;
826 PyObject *it;
827 long start;
828 } dropwhileobject;
829
830 static PyTypeObject dropwhile_type;
831
832 static PyObject *
833 dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
834 {
835 PyObject *func, *seq;
836 PyObject *it;
837 dropwhileobject *lz;
838
839 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
840 return NULL;
841
842 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
843 return NULL;
844
845 /* Get iterator. */
846 it = PyObject_GetIter(seq);
847 if (it == NULL)
848 return NULL;
849
850 /* create dropwhileobject structure */
851 lz = (dropwhileobject *)type->tp_alloc(type, 0);
852 if (lz == NULL) {
853 Py_DECREF(it);
854 return NULL;
855 }
856 Py_INCREF(func);
857 lz->func = func;
858 lz->it = it;
859 lz->start = 0;
860
861 return (PyObject *)lz;
862 }
863
864 static void
865 dropwhile_dealloc(dropwhileobject *lz)
866 {
867 PyObject_GC_UnTrack(lz);
868 Py_XDECREF(lz->func);
869 Py_XDECREF(lz->it);
870 Py_TYPE(lz)->tp_free(lz);
871 }
872
873 static int
874 dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
875 {
876 Py_VISIT(lz->it);
877 Py_VISIT(lz->func);
878 return 0;
879 }
880
881 static PyObject *
882 dropwhile_next(dropwhileobject *lz)
883 {
884 PyObject *item, *good;
885 PyObject *it = lz->it;
886 long ok;
887 PyObject *(*iternext)(PyObject *);
888
889 assert(PyIter_Check(it));
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 assert(PyIter_Check(it));
1035 item = (*Py_TYPE(it)->tp_iternext)(it);
1036 if (item == NULL)
1037 return NULL;
1038
1039 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1040 if (good == NULL) {
1041 Py_DECREF(item);
1042 return NULL;
1043 }
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;
1051 }
1052
1053 PyDoc_STRVAR(takewhile_doc,
1054 "takewhile(predicate, iterable) --> takewhile object\n\
1055 \n\
1056 Return successive entries from an iterable as long as the \n\
1057 predicate evaluates to true for each entry.");
1058
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_compare */
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 */
1101 };
1102
1103
1104 /* islice object ************************************************************/
1105
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;
1114
1115 static PyTypeObject islice_type;
1116
1117 static PyObject *
1118 islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1119 {
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;
1125
1126 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
1127 return NULL;
1128
1129 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1130 return NULL;
1131
1132 numargs = PyTuple_Size(args);
1133 if (numargs == 2) {
1134 if (a1 != Py_None) {
1135 stop = PyInt_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 a non-negative integer or None.");
1141 return NULL;
1142 }
1143 }
1144 } else {
1145 if (a1 != Py_None)
1146 start = PyInt_AsSsize_t(a1);
1147 if (start == -1 && PyErr_Occurred())
1148 PyErr_Clear();
1149 if (a2 != Py_None) {
1150 stop = PyInt_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 a non-negative integer or None.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be non-negative integers or None.");
1163 return NULL;
1164 }
1165
1166 if (a3 != NULL) {
1167 if (a3 != Py_None)
1168 step = PyInt_AsSsize_t(a3);
1169 if (step == -1 && PyErr_Occurred())
1170 PyErr_Clear();
1171 }
1172 if (step<1) {
1173 PyErr_SetString(PyExc_ValueError,
1174 "Step for islice() must be a positive integer or None.");
1175 return NULL;
1176 }
1177
1178 /* Get iterator. */
1179 it = PyObject_GetIter(seq);
1180 if (it == NULL)
1181 return NULL;
1182
1183 /* create isliceobject structure */
1184 lz = (isliceobject *)type->tp_alloc(type, 0);
1185 if (lz == NULL) {
1186 Py_DECREF(it);
1187 return NULL;
1188 }
1189 lz->it = it;
1190 lz->next = start;
1191 lz->stop = stop;
1192 lz->step = step;
1193 lz->cnt = 0L;
1194
1195 return (PyObject *)lz;
1196 }
1197
1198 static void
1199 islice_dealloc(isliceobject *lz)
1200 {
1201 PyObject_GC_UnTrack(lz);
1202 Py_XDECREF(lz->it);
1203 Py_TYPE(lz)->tp_free(lz);
1204 }
1205
1206 static int
1207 islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1208 {
1209 Py_VISIT(lz->it);
1210 return 0;
1211 }
1212
1213 static PyObject *
1214 islice_next(isliceobject *lz)
1215 {
1216 PyObject *item;
1217 PyObject *it = lz->it;
1218 Py_ssize_t oldnext;
1219 PyObject *(*iternext)(PyObject *);
1220
1221 assert(PyIter_Check(it));
1222 iternext = *Py_TYPE(it)->tp_iternext;
1223 while (lz->cnt < lz->next) {
1224 item = iternext(it);
1225 if (item == NULL)
1226 return NULL;
1227 Py_DECREF(item);
1228 lz->cnt++;
1229 }
1230 if (lz->stop != -1 && lz->cnt >= lz->stop)
1231 return NULL;
1232 assert(PyIter_Check(it));
1233 item = iternext(it);
1234 if (item == NULL)
1235 return NULL;
1236 lz->cnt++;
1237 oldnext = lz->next;
1238 lz->next += lz->step;
1239 if (lz->next < oldnext) /* Check for overflow */
1240 lz->next = lz->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 assert(PyIter_Check(it));
1365 args = (*Py_TYPE(it)->tp_iternext)(it);
1366 if (args == NULL)
1367 return NULL;
1368 if (!PyTuple_CheckExact(args)) {
1369 PyObject *newargs = PySequence_Tuple(args);
1370 Py_DECREF(args);
1371 if (newargs == NULL)
1372 return NULL;
1373 args = newargs;
1374 }
1375 result = PyObject_Call(lz->func, args, NULL);
1376 Py_DECREF(args);
1377 return result;
1378 }
1379
1380 PyDoc_STRVAR(starmap_doc,
1381 "starmap(function, sequence) --> starmap object\n\
1382 \n\
1383 Return an iterator whose values are returned from the function evaluated\n\
1384 with a argument tuple taken from the given sequence.");
1385
1386 static PyTypeObject starmap_type = {
1387 PyVarObject_HEAD_INIT(NULL, 0)
1388 "itertools.starmap", /* tp_name */
1389 sizeof(starmapobject), /* tp_basicsize */
1390 0, /* tp_itemsize */
1391 /* methods */
1392 (destructor)starmap_dealloc, /* tp_dealloc */
1393 0, /* tp_print */
1394 0, /* tp_getattr */
1395 0, /* tp_setattr */
1396 0, /* tp_compare */
1397 0, /* tp_repr */
1398 0, /* tp_as_number */
1399 0, /* tp_as_sequence */
1400 0, /* tp_as_mapping */
1401 0, /* tp_hash */
1402 0, /* tp_call */
1403 0, /* tp_str */
1404 PyObject_GenericGetAttr, /* tp_getattro */
1405 0, /* tp_setattro */
1406 0, /* tp_as_buffer */
1407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1408 Py_TPFLAGS_BASETYPE, /* tp_flags */
1409 starmap_doc, /* tp_doc */
1410 (traverseproc)starmap_traverse, /* tp_traverse */
1411 0, /* tp_clear */
1412 0, /* tp_richcompare */
1413 0, /* tp_weaklistoffset */
1414 PyObject_SelfIter, /* tp_iter */
1415 (iternextfunc)starmap_next, /* tp_iternext */
1416 0, /* tp_methods */
1417 0, /* tp_members */
1418 0, /* tp_getset */
1419 0, /* tp_base */
1420 0, /* tp_dict */
1421 0, /* tp_descr_get */
1422 0, /* tp_descr_set */
1423 0, /* tp_dictoffset */
1424 0, /* tp_init */
1425 0, /* tp_alloc */
1426 starmap_new, /* tp_new */
1427 PyObject_GC_Del, /* tp_free */
1428 };
1429
1430
1431 /* imap object ************************************************************/
1432
1433 typedef struct {
1434 PyObject_HEAD
1435 PyObject *iters;
1436 PyObject *func;
1437 } imapobject;
1438
1439 static PyTypeObject imap_type;
1440
1441 static PyObject *
1442 imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1443 {
1444 PyObject *it, *iters, *func;
1445 imapobject *lz;
1446 Py_ssize_t numargs, i;
1447
1448 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1449 return NULL;
1450
1451 numargs = PyTuple_Size(args);
1452 if (numargs < 2) {
1453 PyErr_SetString(PyExc_TypeError,
1454 "imap() must have at least two arguments.");
1455 return NULL;
1456 }
1457
1458 iters = PyTuple_New(numargs-1);
1459 if (iters == NULL)
1460 return NULL;
1461
1462 for (i=1 ; i<numargs ; i++) {
1463 /* Get iterator. */
1464 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1465 if (it == NULL) {
1466 Py_DECREF(iters);
1467 return NULL;
1468 }
1469 PyTuple_SET_ITEM(iters, i-1, it);
1470 }
1471
1472 /* create imapobject structure */
1473 lz = (imapobject *)type->tp_alloc(type, 0);
1474 if (lz == NULL) {
1475 Py_DECREF(iters);
1476 return NULL;
1477 }
1478 lz->iters = iters;
1479 func = PyTuple_GET_ITEM(args, 0);
1480 Py_INCREF(func);
1481 lz->func = func;
1482
1483 return (PyObject *)lz;
1484 }
1485
1486 static void
1487 imap_dealloc(imapobject *lz)
1488 {
1489 PyObject_GC_UnTrack(lz);
1490 Py_XDECREF(lz->iters);
1491 Py_XDECREF(lz->func);
1492 Py_TYPE(lz)->tp_free(lz);
1493 }
1494
1495 static int
1496 imap_traverse(imapobject *lz, visitproc visit, void *arg)
1497 {
1498 Py_VISIT(lz->iters);
1499 Py_VISIT(lz->func);
1500 return 0;
1501 }
1502
1503 /*
1504 imap() is an iterator version of __builtins__.map() except that it does
1505 not have the None fill-in feature. That was intentionally left out for
1506 the following reasons:
1507
1508 1) Itertools are designed to be easily combined and chained together.
1509 Having all tools stop with the shortest input is a unifying principle
1510 that makes it easier to combine finite iterators (supplying data) with
1511 infinite iterators like count() and repeat() (for supplying sequential
1512 or constant arguments to a function).
1513
1514 2) In typical use cases for combining itertools, having one finite data
1515 supplier run out before another is likely to be an error condition which
1516 should not pass silently by automatically supplying None.
1517
1518 3) The use cases for automatic None fill-in are rare -- not many functions
1519 do something useful when a parameter suddenly switches type and becomes
1520 None.
1521
1522 4) If a need does arise, it can be met by __builtins__.map() or by
1523 writing: chain(iterable, repeat(None)).
1524
1525 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1526 */
1527
1528 static PyObject *
1529 imap_next(imapobject *lz)
1530 {
1531 PyObject *val;
1532 PyObject *argtuple;
1533 PyObject *result;
1534 Py_ssize_t numargs, i;
1535
1536 numargs = PyTuple_Size(lz->iters);
1537 argtuple = PyTuple_New(numargs);
1538 if (argtuple == NULL)
1539 return NULL;
1540
1541 for (i=0 ; i<numargs ; i++) {
1542 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1543 if (val == NULL) {
1544 Py_DECREF(argtuple);
1545 return NULL;
1546 }
1547 PyTuple_SET_ITEM(argtuple, i, val);
1548 }
1549 if (lz->func == Py_None)
1550 return argtuple;
1551 result = PyObject_Call(lz->func, argtuple, NULL);
1552 Py_DECREF(argtuple);
1553 return result;
1554 }
1555
1556 PyDoc_STRVAR(imap_doc,
1557 "imap(func, *iterables) --> imap object\n\
1558 \n\
1559 Make an iterator that computes the function using arguments from\n\
1560 each of the iterables. Like map() except that it returns\n\
1561 an iterator instead of a list and that it stops when the shortest\n\
1562 iterable is exhausted instead of filling in None for shorter\n\
1563 iterables.");
1564
1565 static PyTypeObject imap_type = {
1566 PyVarObject_HEAD_INIT(NULL, 0)
1567 "itertools.imap", /* tp_name */
1568 sizeof(imapobject), /* tp_basicsize */
1569 0, /* tp_itemsize */
1570 /* methods */
1571 (destructor)imap_dealloc, /* tp_dealloc */
1572 0, /* tp_print */
1573 0, /* tp_getattr */
1574 0, /* tp_setattr */
1575 0, /* tp_compare */
1576 0, /* tp_repr */
1577 0, /* tp_as_number */
1578 0, /* tp_as_sequence */
1579 0, /* tp_as_mapping */
1580 0, /* tp_hash */
1581 0, /* tp_call */
1582 0, /* tp_str */
1583 PyObject_GenericGetAttr, /* tp_getattro */
1584 0, /* tp_setattro */
1585 0, /* tp_as_buffer */
1586 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1587 Py_TPFLAGS_BASETYPE, /* tp_flags */
1588 imap_doc, /* tp_doc */
1589 (traverseproc)imap_traverse, /* tp_traverse */
1590 0, /* tp_clear */
1591 0, /* tp_richcompare */
1592 0, /* tp_weaklistoffset */
1593 PyObject_SelfIter, /* tp_iter */
1594 (iternextfunc)imap_next, /* tp_iternext */
1595 0, /* tp_methods */
1596 0, /* tp_members */
1597 0, /* tp_getset */
1598 0, /* tp_base */
1599 0, /* tp_dict */
1600 0, /* tp_descr_get */
1601 0, /* tp_descr_set */
1602 0, /* tp_dictoffset */
1603 0, /* tp_init */
1604 0, /* tp_alloc */
1605 imap_new, /* tp_new */
1606 PyObject_GC_Del, /* tp_free */
1607 };
1608
1609
1610 /* chain object ************************************************************/
1611
1612 typedef struct {
1613 PyObject_HEAD
1614 PyObject *source; /* Iterator over input iterables */
1615 PyObject *active; /* Currently running input iterator */
1616 } chainobject;
1617
1618 static PyTypeObject chain_type;
1619
1620 static PyObject *
1621 chain_new_internal(PyTypeObject *type, PyObject *source)
1622 {
1623 chainobject *lz;
1624
1625 lz = (chainobject *)type->tp_alloc(type, 0);
1626 if (lz == NULL) {
1627 Py_DECREF(source);
1628 return NULL;
1629 }
1630
1631 lz->source = source;
1632 lz->active = NULL;
1633 return (PyObject *)lz;
1634 }
1635
1636 static PyObject *
1637 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1638 {
1639 PyObject *source;
1640
1641 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1642 return NULL;
1643
1644 source = PyObject_GetIter(args);
1645 if (source == NULL)
1646 return NULL;
1647
1648 return chain_new_internal(type, source);
1649 }
1650
1651 static PyObject *
1652 chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1653 {
1654 PyObject *source;
1655
1656 source = PyObject_GetIter(arg);
1657 if (source == NULL)
1658 return NULL;
1659
1660 return chain_new_internal(type, source);
1661 }
1662
1663 static void
1664 chain_dealloc(chainobject *lz)
1665 {
1666 PyObject_GC_UnTrack(lz);
1667 Py_XDECREF(lz->active);
1668 Py_XDECREF(lz->source);
1669 Py_TYPE(lz)->tp_free(lz);
1670 }
1671
1672 static int
1673 chain_traverse(chainobject *lz, visitproc visit, void *arg)
1674 {
1675 Py_VISIT(lz->source);
1676 Py_VISIT(lz->active);
1677 return 0;
1678 }
1679
1680 static PyObject *
1681 chain_next(chainobject *lz)
1682 {
1683 PyObject *item;
1684
1685 if (lz->source == NULL)
1686 return NULL; /* already stopped */
1687
1688 if (lz->active == NULL) {
1689 PyObject *iterable = PyIter_Next(lz->source);
1690 if (iterable == NULL) {
1691 Py_CLEAR(lz->source);
1692 return NULL; /* no more input sources */
1693 }
1694 lz->active = PyObject_GetIter(iterable);
1695 Py_DECREF(iterable);
1696 if (lz->active == NULL) {
1697 Py_CLEAR(lz->source);
1698 return NULL; /* input not iterable */
1699 }
1700 }
1701 item = PyIter_Next(lz->active);
1702 if (item != NULL)
1703 return item;
1704 if (PyErr_Occurred()) {
1705 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1706 PyErr_Clear();
1707 else
1708 return NULL; /* input raised an exception */
1709 }
1710 Py_CLEAR(lz->active);
1711 return chain_next(lz); /* recurse and use next active */
1712 }
1713
1714 PyDoc_STRVAR(chain_doc,
1715 "chain(*iterables) --> chain object\n\
1716 \n\
1717 Return a chain object whose .next() method returns elements from the\n\
1718 first iterable until it is exhausted, then elements from the next\n\
1719 iterable, until all of the iterables are exhausted.");
1720
1721 PyDoc_STRVAR(chain_from_iterable_doc,
1722 "chain.from_iterable(iterable) --> chain object\n\
1723 \n\
1724 Alternate chain() contructor taking a single iterable argument\n\
1725 that evaluates lazily.");
1726
1727 static PyMethodDef chain_methods[] = {
1728 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1729 chain_from_iterable_doc},
1730 {NULL, NULL} /* sentinel */
1731 };
1732
1733 static PyTypeObject chain_type = {
1734 PyVarObject_HEAD_INIT(NULL, 0)
1735 "itertools.chain", /* tp_name */
1736 sizeof(chainobject), /* tp_basicsize */
1737 0, /* tp_itemsize */
1738 /* methods */
1739 (destructor)chain_dealloc, /* tp_dealloc */
1740 0, /* tp_print */
1741 0, /* tp_getattr */
1742 0, /* tp_setattr */
1743 0, /* tp_compare */
1744 0, /* tp_repr */
1745 0, /* tp_as_number */
1746 0, /* tp_as_sequence */
1747 0, /* tp_as_mapping */
1748 0, /* tp_hash */
1749 0, /* tp_call */
1750 0, /* tp_str */
1751 PyObject_GenericGetAttr, /* tp_getattro */
1752 0, /* tp_setattro */
1753 0, /* tp_as_buffer */
1754 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1755 Py_TPFLAGS_BASETYPE, /* tp_flags */
1756 chain_doc, /* tp_doc */
1757 (traverseproc)chain_traverse, /* tp_traverse */
1758 0, /* tp_clear */
1759 0, /* tp_richcompare */
1760 0, /* tp_weaklistoffset */
1761 PyObject_SelfIter, /* tp_iter */
1762 (iternextfunc)chain_next, /* tp_iternext */
1763 chain_methods, /* tp_methods */
1764 0, /* tp_members */
1765 0, /* tp_getset */
1766 0, /* tp_base */
1767 0, /* tp_dict */
1768 0, /* tp_descr_get */
1769 0, /* tp_descr_set */
1770 0, /* tp_dictoffset */
1771 0, /* tp_init */
1772 0, /* tp_alloc */
1773 chain_new, /* tp_new */
1774 PyObject_GC_Del, /* tp_free */
1775 };
1776
1777
1778 /* product object ************************************************************/
1779
1780 typedef struct {
1781 PyObject_HEAD
1782 PyObject *pools; /* tuple of pool tuples */
1783 Py_ssize_t *indices; /* one index per pool */
1784 PyObject *result; /* most recently returned result tuple */
1785 int stopped; /* set to 1 when the product iterator is exhausted */
1786 } productobject;
1787
1788 static PyTypeObject product_type;
1789
1790 static PyObject *
1791 product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1792 {
1793 productobject *lz;
1794 Py_ssize_t nargs, npools, repeat=1;
1795 PyObject *pools = NULL;
1796 Py_ssize_t *indices = NULL;
1797 Py_ssize_t i;
1798
1799 if (kwds != NULL) {
1800 char *kwlist[] = {"repeat", 0};
1801 PyObject *tmpargs = PyTuple_New(0);
1802 if (tmpargs == NULL)
1803 return NULL;
1804 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1805 Py_DECREF(tmpargs);
1806 return NULL;
1807 }
1808 Py_DECREF(tmpargs);
1809 if (repeat < 0) {
1810 PyErr_SetString(PyExc_ValueError,
1811 "repeat argument cannot be negative");
1812 return NULL;
1813 }
1814 }
1815
1816 assert(PyTuple_Check(args));
1817 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1818 npools = nargs * repeat;
1819
1820 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1821 if (indices == NULL) {
1822 PyErr_NoMemory();
1823 goto error;
1824 }
1825
1826 pools = PyTuple_New(npools);
1827 if (pools == NULL)
1828 goto error;
1829
1830 for (i=0; i < nargs ; ++i) {
1831 PyObject *item = PyTuple_GET_ITEM(args, i);
1832 PyObject *pool = PySequence_Tuple(item);
1833 if (pool == NULL)
1834 goto error;
1835 PyTuple_SET_ITEM(pools, i, pool);
1836 indices[i] = 0;
1837 }
1838 for ( ; i < npools; ++i) {
1839 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1840 Py_INCREF(pool);
1841 PyTuple_SET_ITEM(pools, i, pool);
1842 indices[i] = 0;
1843 }
1844
1845 /* create productobject structure */
1846 lz = (productobject *)type->tp_alloc(type, 0);
1847 if (lz == NULL)
1848 goto error;
1849
1850 lz->pools = pools;
1851 lz->indices = indices;
1852 lz->result = NULL;
1853 lz->stopped = 0;
1854
1855 return (PyObject *)lz;
1856
1857 error:
1858 if (indices != NULL)
1859 PyMem_Free(indices);
1860 Py_XDECREF(pools);
1861 return NULL;
1862 }
1863
1864 static void
1865 product_dealloc(productobject *lz)
1866 {
1867 PyObject_GC_UnTrack(lz);
1868 Py_XDECREF(lz->pools);
1869 Py_XDECREF(lz->result);
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 product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1978 product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1979
1980 static PyTypeObject product_type = {
1981 PyVarObject_HEAD_INIT(NULL, 0)
1982 "itertools.product", /* tp_name */
1983 sizeof(productobject), /* tp_basicsize */
1984 0, /* tp_itemsize */
1985 /* methods */
1986 (destructor)product_dealloc, /* tp_dealloc */
1987 0, /* tp_print */
1988 0, /* tp_getattr */
1989 0, /* tp_setattr */
1990 0, /* tp_compare */
1991 0, /* tp_repr */
1992 0, /* tp_as_number */
1993 0, /* tp_as_sequence */
1994 0, /* tp_as_mapping */
1995 0, /* tp_hash */
1996 0, /* tp_call */
1997 0, /* tp_str */
1998 PyObject_GenericGetAttr, /* tp_getattro */
1999 0, /* tp_setattro */
2000 0, /* tp_as_buffer */
2001 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2002 Py_TPFLAGS_BASETYPE, /* tp_flags */
2003 product_doc, /* tp_doc */
2004 (traverseproc)product_traverse, /* tp_traverse */
2005 0, /* tp_clear */
2006 0, /* tp_richcompare */
2007 0, /* tp_weaklistoffset */
2008 PyObject_SelfIter, /* tp_iter */
2009 (iternextfunc)product_next, /* tp_iternext */
2010 0, /* tp_methods */
2011 0, /* tp_members */
2012 0, /* tp_getset */
2013 0, /* tp_base */
2014 0, /* tp_dict */
2015 0, /* tp_descr_get */
2016 0, /* tp_descr_set */
2017 0, /* tp_dictoffset */
2018 0, /* tp_init */
2019 0, /* tp_alloc */
2020 product_new, /* tp_new */
2021 PyObject_GC_Del, /* tp_free */
2022 };
2023
2024
2025 /* combinations object ************************************************************/
2026
2027 typedef struct {
2028 PyObject_HEAD
2029 PyObject *pool; /* input converted to a tuple */
2030 Py_ssize_t *indices; /* one index per result element */
2031 PyObject *result; /* most recently returned result tuple */
2032 Py_ssize_t r; /* size of result tuple */
2033 int stopped; /* set to 1 when the combinations iterator is exhausted */
2034 } combinationsobject;
2035
2036 static PyTypeObject combinations_type;
2037
2038 static PyObject *
2039 combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2040 {
2041 combinationsobject *co;
2042 Py_ssize_t n;
2043 Py_ssize_t r;
2044 PyObject *pool = NULL;
2045 PyObject *iterable = NULL;
2046 Py_ssize_t *indices = NULL;
2047 Py_ssize_t i;
2048 static char *kwargs[] = {"iterable", "r", NULL};
2049
2050 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2051 &iterable, &r))
2052 return NULL;
2053
2054 pool = PySequence_Tuple(iterable);
2055 if (pool == NULL)
2056 goto error;
2057 n = PyTuple_GET_SIZE(pool);
2058 if (r < 0) {
2059 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2060 goto error;
2061 }
2062 if (r > n) {
2063 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
2064 goto error;
2065 }
2066
2067 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2068 if (indices == NULL) {
2069 PyErr_NoMemory();
2070 goto error;
2071 }
2072
2073 for (i=0 ; i<r ; i++)
2074 indices[i] = i;
2075
2076 /* create combinationsobject structure */
2077 co = (combinationsobject *)type->tp_alloc(type, 0);
2078 if (co == NULL)
2079 goto error;
2080
2081 co->pool = pool;
2082 co->indices = indices;
2083 co->result = NULL;
2084 co->r = r;
2085 co->stopped = 0;
2086
2087 return (PyObject *)co;
2088
2089 error:
2090 if (indices != NULL)
2091 PyMem_Free(indices);
2092 Py_XDECREF(pool);
2093 return NULL;
2094 }
2095
2096 static void
2097 combinations_dealloc(combinationsobject *co)
2098 {
2099 PyObject_GC_UnTrack(co);
2100 Py_XDECREF(co->pool);
2101 Py_XDECREF(co->result);
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 /* permutations object ************************************************************
2252
2253 def permutations(iterable, r=None):
2254 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2255 pool = tuple(iterable)
2256 n = len(pool)
2257 r = n if r is None else r
2258 indices = range(n)
2259 cycles = range(n-r+1, n+1)[::-1]
2260 yield tuple(pool[i] for i in indices[:r])
2261 while n:
2262 for i in reversed(range(r)):
2263 cycles[i] -= 1
2264 if cycles[i] == 0:
2265 indices[i:] = indices[i+1:] + indices[i:i+1]
2266 cycles[i] = n - i
2267 else:
2268 j = cycles[i]
2269 indices[i], indices[-j] = indices[-j], indices[i]
2270 yield tuple(pool[i] for i in indices[:r])
2271 break
2272 else:
2273 return
2274 */
2275
2276 typedef struct {
2277 PyObject_HEAD
2278 PyObject *pool; /* input converted to a tuple */
2279 Py_ssize_t *indices; /* one index per element in the pool */
2280 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2281 PyObject *result; /* most recently returned result tuple */
2282 Py_ssize_t r; /* size of result tuple */
2283 int stopped; /* set to 1 when the permutations iterator is exhausted */
2284 } permutationsobject;
2285
2286 static PyTypeObject permutations_type;
2287
2288 static PyObject *
2289 permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2290 {
2291 permutationsobject *po;
2292 Py_ssize_t n;
2293 Py_ssize_t r;
2294 PyObject *robj = Py_None;
2295 PyObject *pool = NULL;
2296 PyObject *iterable = NULL;
2297 Py_ssize_t *indices = NULL;
2298 Py_ssize_t *cycles = NULL;
2299 Py_ssize_t i;
2300 static char *kwargs[] = {"iterable", "r", NULL};
2301
2302 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2303 &iterable, &robj))
2304 return NULL;
2305
2306 pool = PySequence_Tuple(iterable);
2307 if (pool == NULL)
2308 goto error;
2309 n = PyTuple_GET_SIZE(pool);
2310
2311 r = n;
2312 if (robj != Py_None) {
2313 r = PyInt_AsSsize_t(robj);
2314 if (r == -1 && PyErr_Occurred())
2315 goto error;
2316 }
2317 if (r < 0) {
2318 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2319 goto error;
2320 }
2321 if (r > n) {
2322 PyErr_SetString(PyExc_ValueError, "r cannot be bigger than the iterable");
2323 goto error;
2324 }
2325
2326 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2327 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2328 if (indices == NULL || cycles == NULL) {
2329 PyErr_NoMemory();
2330 goto error;
2331 }
2332
2333 for (i=0 ; i<n ; i++)
2334 indices[i] = i;
2335 for (i=0 ; i<r ; i++)
2336 cycles[i] = n - i;
2337
2338 /* create permutationsobject structure */
2339 po = (permutationsobject *)type->tp_alloc(type, 0);
2340 if (po == NULL)
2341 goto error;
2342
2343 po->pool = pool;
2344 po->indices = indices;
2345 po->cycles = cycles;
2346 po->result = NULL;
2347 po->r = r;
2348 po->stopped = 0;
2349
2350 return (PyObject *)po;
2351
2352 error:
2353 if (indices != NULL)
2354 PyMem_Free(indices);
2355 if (cycles != NULL)
2356 PyMem_Free(cycles);
2357 Py_XDECREF(pool);
2358 return NULL;
2359 }
2360
2361 static void
2362 permutations_dealloc(permutationsobject *po)
2363 {
2364 PyObject_GC_UnTrack(po);
2365 Py_XDECREF(po->pool);
2366 Py_XDECREF(po->result);
2367 PyMem_Free(po->indices);
2368 PyMem_Free(po->cycles);
2369 Py_TYPE(po)->tp_free(po);
2370 }
2371
2372 static int
2373 permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2374 {
2375 Py_VISIT(po->pool);
2376 Py_VISIT(po->result);
2377 return 0;
2378 }
2379
2380 static PyObject *
2381 permutations_next(permutationsobject *po)
2382 {
2383 PyObject *elem;
2384 PyObject *oldelem;
2385 PyObject *pool = po->pool;
2386 Py_ssize_t *indices = po->indices;
2387 Py_ssize_t *cycles = po->cycles;
2388 PyObject *result = po->result;
2389 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2390 Py_ssize_t r = po->r;
2391 Py_ssize_t i, j, k, index;
2392
2393 if (po->stopped)
2394 return NULL;
2395
2396 if (result == NULL) {
2397 /* On the first pass, initialize result tuple using the indices */
2398 result = PyTuple_New(r);
2399 if (result == NULL)
2400 goto empty;
2401 po->result = result;
2402 for (i=0; i<r ; i++) {
2403 index = indices[i];
2404 elem = PyTuple_GET_ITEM(pool, index);
2405 Py_INCREF(elem);
2406 PyTuple_SET_ITEM(result, i, elem);
2407 }
2408 } else {
2409 if (n == 0)
2410 goto empty;
2411
2412 /* Copy the previous result tuple or re-use it if available */
2413 if (Py_REFCNT(result) > 1) {
2414 PyObject *old_result = result;
2415 result = PyTuple_New(r);
2416 if (result == NULL)
2417 goto empty;
2418 po->result = result;
2419 for (i=0; i<r ; i++) {
2420 elem = PyTuple_GET_ITEM(old_result, i);
2421 Py_INCREF(elem);
2422 PyTuple_SET_ITEM(result, i, elem);
2423 }
2424 Py_DECREF(old_result);
2425 }
2426 /* Now, we've got the only copy so we can update it in-place */
2427 assert(r == 0 || Py_REFCNT(result) == 1);
2428
2429 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2430 for (i=r-1 ; i>=0 ; i--) {
2431 cycles[i] -= 1;
2432 if (cycles[i] == 0) {
2433 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2434 index = indices[i];
2435 for (j=i ; j<n-1 ; j++)
2436 indices[j] = indices[j+1];
2437 indices[n-1] = index;
2438 cycles[i] = n - i;
2439 } else {
2440 j = cycles[i];
2441 index = indices[i];
2442 indices[i] = indices[n-j];
2443 indices[n-j] = index;
2444
2445 for (k=i; k<r ; k++) {
2446 /* start with i, the leftmost element that changed */
2447 /* yield tuple(pool[k] for k in indices[:r]) */
2448 index = indices[k];
2449 elem = PyTuple_GET_ITEM(pool, index);
2450 Py_INCREF(elem);
2451 oldelem = PyTuple_GET_ITEM(result, k);
2452 PyTuple_SET_ITEM(result, k, elem);
2453 Py_DECREF(oldelem);
2454 }
2455 break;
2456 }
2457 }
2458 /* If i is negative, then the cycles have all
2459 rolled-over and we're done. */
2460 if (i < 0)
2461 goto empty;
2462 }
2463 Py_INCREF(result);
2464 return result;
2465
2466 empty:
2467 po->stopped = 1;
2468 return NULL;
2469 }
2470
2471 PyDoc_STRVAR(permutations_doc,
2472 "permutations(iterable[, r]) --> permutations object\n\
2473 \n\
2474 Return successive r-length permutations of elements in the iterable.\n\n\
2475 permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
2476
2477 static PyTypeObject permutations_type = {
2478 PyVarObject_HEAD_INIT(NULL, 0)
2479 "itertools.permutations", /* tp_name */
2480 sizeof(permutationsobject), /* tp_basicsize */
2481 0, /* tp_itemsize */
2482 /* methods */
2483 (destructor)permutations_dealloc, /* tp_dealloc */
2484 0, /* tp_print */
2485 0, /* tp_getattr */
2486 0, /* tp_setattr */
2487 0, /* tp_compare */
2488 0, /* tp_repr */
2489 0, /* tp_as_number */
2490 0, /* tp_as_sequence */
2491 0, /* tp_as_mapping */
2492 0, /* tp_hash */
2493 0, /* tp_call */
2494 0, /* tp_str */
2495 PyObject_GenericGetAttr, /* tp_getattro */
2496 0, /* tp_setattro */
2497 0, /* tp_as_buffer */
2498 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2499 Py_TPFLAGS_BASETYPE, /* tp_flags */
2500 permutations_doc, /* tp_doc */
2501 (traverseproc)permutations_traverse, /* tp_traverse */
2502 0, /* tp_clear */
2503 0, /* tp_richcompare */
2504 0, /* tp_weaklistoffset */
2505 PyObject_SelfIter, /* tp_iter */
2506 (iternextfunc)permutations_next, /* tp_iternext */
2507 0, /* tp_methods */
2508 0, /* tp_members */
2509 0, /* tp_getset */
2510 0, /* tp_base */
2511 0, /* tp_dict */
2512 0, /* tp_descr_get */
2513 0, /* tp_descr_set */
2514 0, /* tp_dictoffset */
2515 0, /* tp_init */
2516 0, /* tp_alloc */
2517 permutations_new, /* tp_new */
2518 PyObject_GC_Del, /* tp_free */
2519 };
2520
2521
2522 /* ifilter object ************************************************************/
2523
2524 typedef struct {
2525 PyObject_HEAD
2526 PyObject *func;
2527 PyObject *it;
2528 } ifilterobject;
2529
2530 static PyTypeObject ifilter_type;
2531
2532 static PyObject *
2533 ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2534 {
2535 PyObject *func, *seq;
2536 PyObject *it;
2537 ifilterobject *lz;
2538
2539 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2540 return NULL;
2541
2542 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2543 return NULL;
2544
2545 /* Get iterator. */
2546 it = PyObject_GetIter(seq);
2547 if (it == NULL)
2548 return NULL;
2549
2550 /* create ifilterobject structure */
2551 lz = (ifilterobject *)type->tp_alloc(type, 0);
2552 if (lz == NULL) {
2553 Py_DECREF(it);
2554 return NULL;
2555 }
2556 Py_INCREF(func);
2557 lz->func = func;
2558 lz->it = it;
2559
2560 return (PyObject *)lz;
2561 }
2562
2563 static void
2564 ifilter_dealloc(ifilterobject *lz)
2565 {
2566 PyObject_GC_UnTrack(lz);
2567 Py_XDECREF(lz->func);
2568 Py_XDECREF(lz->it);
2569 Py_TYPE(lz)->tp_free(lz);
2570 }
2571
2572 static int
2573 ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2574 {
2575 Py_VISIT(lz->it);
2576 Py_VISIT(lz->func);
2577 return 0;
2578 }
2579
2580 static PyObject *
2581 ifilter_next(ifilterobject *lz)
2582 {
2583 PyObject *item;
2584 PyObject *it = lz->it;
2585 long ok;
2586 PyObject *(*iternext)(PyObject *);
2587
2588 assert(PyIter_Check(it));
2589 iternext = *Py_TYPE(it)->tp_iternext;
2590 for (;;) {
2591 item = iternext(it);
2592 if (item == NULL)
2593 return NULL;
2594
2595 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2596 ok = PyObject_IsTrue(item);
2597 } else {
2598 PyObject *good;
2599 good = PyObject_CallFunctionObjArgs(lz->func,
2600 item, NULL);
2601 if (good == NULL) {
2602 Py_DECREF(item);
2603 return NULL;
2604 }
2605 ok = PyObject_IsTrue(good);
2606 Py_DECREF(good);
2607 }
2608 if (ok)
2609 return item;
2610 Py_DECREF(item);
2611 }
2612 }
2613
2614 PyDoc_STRVAR(ifilter_doc,
2615 "ifilter(function or None, sequence) --> ifilter object\n\
2616 \n\
2617 Return those items of sequence for which function(item) is true.\n\
2618 If function is None, return the items that are true.");
2619
2620 static PyTypeObject ifilter_type = {
2621 PyVarObject_HEAD_INIT(NULL, 0)
2622 "itertools.ifilter", /* tp_name */
2623 sizeof(ifilterobject), /* tp_basicsize */
2624 0, /* tp_itemsize */
2625 /* methods */
2626 (destructor)ifilter_dealloc, /* tp_dealloc */
2627 0, /* tp_print */
2628 0, /* tp_getattr */
2629 0, /* tp_setattr */
2630 0, /* tp_compare */
2631 0, /* tp_repr */
2632 0, /* tp_as_number */
2633 0, /* tp_as_sequence */
2634 0, /* tp_as_mapping */
2635 0, /* tp_hash */
2636 0, /* tp_call */
2637 0, /* tp_str */
2638 PyObject_GenericGetAttr, /* tp_getattro */
2639 0, /* tp_setattro */
2640 0, /* tp_as_buffer */
2641 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2642 Py_TPFLAGS_BASETYPE, /* tp_flags */
2643 ifilter_doc, /* tp_doc */
2644 (traverseproc)ifilter_traverse, /* tp_traverse */
2645 0, /* tp_clear */
2646 0, /* tp_richcompare */
2647 0, /* tp_weaklistoffset */
2648 PyObject_SelfIter, /* tp_iter */
2649 (iternextfunc)ifilter_next, /* tp_iternext */
2650 0, /* tp_methods */
2651 0, /* tp_members */
2652 0, /* tp_getset */
2653 0, /* tp_base */
2654 0, /* tp_dict */
2655 0, /* tp_descr_get */
2656 0, /* tp_descr_set */
2657 0, /* tp_dictoffset */
2658 0, /* tp_init */
2659 0, /* tp_alloc */
2660 ifilter_new, /* tp_new */
2661 PyObject_GC_Del, /* tp_free */
2662 };
2663
2664
2665 /* ifilterfalse object ************************************************************/
2666
2667 typedef struct {
2668 PyObject_HEAD
2669 PyObject *func;
2670 PyObject *it;
2671 } ifilterfalseobject;
2672
2673 static PyTypeObject ifilterfalse_type;
2674
2675 static PyObject *
2676 ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2677 {
2678 PyObject *func, *seq;
2679 PyObject *it;
2680 ifilterfalseobject *lz;
2681
2682 if (type == &ifilterfalse_type &&
2683 !_PyArg_NoKeywords("ifilterfalse()", kwds))
2684 return NULL;
2685
2686 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2687 return NULL;
2688
2689 /* Get iterator. */
2690 it = PyObject_GetIter(seq);
2691 if (it == NULL)
2692 return NULL;
2693
2694 /* create ifilterfalseobject structure */
2695 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
2696 if (lz == NULL) {
2697 Py_DECREF(it);
2698 return NULL;
2699 }
2700 Py_INCREF(func);
2701 lz->func = func;
2702 lz->it = it;
2703
2704 return (PyObject *)lz;
2705 }
2706
2707 static void
2708 ifilterfalse_dealloc(ifilterfalseobject *lz)
2709 {
2710 PyObject_GC_UnTrack(lz);
2711 Py_XDECREF(lz->func);
2712 Py_XDECREF(lz->it);
2713 Py_TYPE(lz)->tp_free(lz);
2714 }
2715
2716 static int
2717 ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
2718 {
2719 Py_VISIT(lz->it);
2720 Py_VISIT(lz->func);
2721 return 0;
2722 }
2723
2724 static PyObject *
2725 ifilterfalse_next(ifilterfalseobject *lz)
2726 {
2727 PyObject *item;
2728 PyObject *it = lz->it;
2729 long ok;
2730 PyObject *(*iternext)(PyObject *);
2731
2732 assert(PyIter_Check(it));
2733 iternext = *Py_TYPE(it)->tp_iternext;
2734 for (;;) {
2735 item = iternext(it);
2736 if (item == NULL)
2737 return NULL;
2738
2739 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2740 ok = PyObject_IsTrue(item);
2741 } else {
2742 PyObject *good;
2743 good = PyObject_CallFunctionObjArgs(lz->func,
2744 item, NULL);
2745 if (good == NULL) {
2746 Py_DECREF(item);
2747 return NULL;
2748 }
2749 ok = PyObject_IsTrue(good);
2750 Py_DECREF(good);
2751 }
2752 if (!ok)
2753 return item;
2754 Py_DECREF(item);
2755 }
2756 }
2757
2758 PyDoc_STRVAR(ifilterfalse_doc,
2759 "ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2760 \n\
2761 Return those items of sequence for which function(item) is false.\n\
2762 If function is None, return the items that are false.");
2763
2764 static PyTypeObject ifilterfalse_type = {
2765 PyVarObject_HEAD_INIT(NULL, 0)
2766 "itertools.ifilterfalse", /* tp_name */
2767 sizeof(ifilterfalseobject), /* tp_basicsize */
2768 0, /* tp_itemsize */
2769 /* methods */
2770 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2771 0, /* tp_print */
2772 0, /* tp_getattr */
2773 0, /* tp_setattr */
2774 0, /* tp_compare */
2775 0, /* tp_repr */
2776 0, /* tp_as_number */
2777 0, /* tp_as_sequence */
2778 0, /* tp_as_mapping */
2779 0, /* tp_hash */
2780 0, /* tp_call */
2781 0, /* tp_str */
2782 PyObject_GenericGetAttr, /* tp_getattro */
2783 0, /* tp_setattro */
2784 0, /* tp_as_buffer */
2785 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2786 Py_TPFLAGS_BASETYPE, /* tp_flags */
2787 ifilterfalse_doc, /* tp_doc */
2788 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2789 0, /* tp_clear */
2790 0, /* tp_richcompare */
2791 0, /* tp_weaklistoffset */
2792 PyObject_SelfIter, /* tp_iter */
2793 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2794 0, /* tp_methods */
2795 0, /* tp_members */
2796 0, /* tp_getset */
2797 0, /* tp_base */
2798 0, /* tp_dict */
2799 0, /* tp_descr_get */
2800 0, /* tp_descr_set */
2801 0, /* tp_dictoffset */
2802 0, /* tp_init */
2803 0, /* tp_alloc */
2804 ifilterfalse_new, /* tp_new */
2805 PyObject_GC_Del, /* tp_free */
2806 };
2807
2808
2809 /* count object ************************************************************/
2810
2811 typedef struct {
2812 PyObject_HEAD
2813 Py_ssize_t cnt;
2814 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
2815 } countobject;
2816
2817 static PyTypeObject count_type;
2818
2819 static PyObject *
2820 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2821 {
2822 countobject *lz;
2823 Py_ssize_t cnt = 0;
2824 PyObject *cnt_arg = NULL;
2825 PyObject *long_cnt = NULL;
2826
2827 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
2828 return NULL;
2829
2830 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
2831 return NULL;
2832
2833 if (cnt_arg != NULL) {
2834 cnt = PyInt_AsSsize_t(cnt_arg);
2835 if (cnt == -1 && PyErr_Occurred()) {
2836 PyErr_Clear();
2837 if (!PyLong_Check(cnt_arg)) {
2838 PyErr_SetString(PyExc_TypeError, "an integer is required");
2839 return NULL;
2840 }
2841 long_cnt = cnt_arg;
2842 Py_INCREF(long_cnt);
2843 cnt = PY_SSIZE_T_MAX;
2844 }
2845 }
2846
2847 /* create countobject structure */
2848 lz = (countobject *)PyObject_New(countobject, &count_type);
2849 if (lz == NULL) {
2850 Py_XDECREF(long_cnt);
2851 return NULL;
2852 }
2853 lz->cnt = cnt;
2854 lz->long_cnt = long_cnt;
2855
2856 return (PyObject *)lz;
2857 }
2858
2859 static void
2860 count_dealloc(countobject *lz)
2861 {
2862 Py_XDECREF(lz->long_cnt);
2863 PyObject_Del(lz);
2864 }
2865
2866 static PyObject *
2867 count_nextlong(countobject *lz)
2868 {
2869 static PyObject *one = NULL;
2870 PyObject *cnt;
2871 PyObject *stepped_up;
2872
2873 if (lz->long_cnt == NULL) {
2874 lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
2875 if (lz->long_cnt == NULL)
2876 return NULL;
2877 }
2878 if (one == NULL) {
2879 one = PyInt_FromLong(1);
2880 if (one == NULL)
2881 return NULL;
2882 }
2883 cnt = lz->long_cnt;
2884 assert(cnt != NULL);
2885 stepped_up = PyNumber_Add(cnt, one);
2886 if (stepped_up == NULL)
2887 return NULL;
2888 lz->long_cnt = stepped_up;
2889 return cnt;
2890 }
2891
2892 static PyObject *
2893 count_next(countobject *lz)
2894 {
2895 if (lz->cnt == PY_SSIZE_T_MAX)
2896 return count_nextlong(lz);
2897 return PyInt_FromSsize_t(lz->cnt++);
2898 }
2899
2900 static PyObject *
2901 count_repr(countobject *lz)
2902 {
2903 PyObject *cnt_repr;
2904 PyObject *result;
2905
2906 if (lz->cnt != PY_SSIZE_T_MAX)
2907 return PyString_FromFormat("count(%zd)", lz->cnt);
2908
2909 cnt_repr = PyObject_Repr(lz->long_cnt);
2910 if (cnt_repr == NULL)
2911 return NULL;
2912 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
2913 Py_DECREF(cnt_repr);
2914 return result;
2915 }
2916
2917 PyDoc_STRVAR(count_doc,
2918 "count([firstval]) --> count object\n\
2919 \n\
2920 Return a count object whose .next() method returns consecutive\n\
2921 integers starting from zero or, if specified, from firstval.");
2922
2923 static PyTypeObject count_type = {
2924 PyVarObject_HEAD_INIT(NULL, 0)
2925 "itertools.count", /* tp_name */
2926 sizeof(countobject), /* tp_basicsize */
2927 0, /* tp_itemsize */
2928 /* methods */
2929 (destructor)count_dealloc, /* tp_dealloc */
2930 0, /* tp_print */
2931 0, /* tp_getattr */
2932 0, /* tp_setattr */
2933 0, /* tp_compare */
2934 (reprfunc)count_repr, /* tp_repr */
2935 0, /* tp_as_number */
2936 0, /* tp_as_sequence */
2937 0, /* tp_as_mapping */
2938 0, /* tp_hash */
2939 0, /* tp_call */
2940 0, /* tp_str */
2941 PyObject_GenericGetAttr, /* tp_getattro */
2942 0, /* tp_setattro */
2943 0, /* tp_as_buffer */
2944 Py_TPFLAGS_DEFAULT, /* tp_flags */
2945 count_doc, /* tp_doc */
2946 0, /* tp_traverse */
2947 0, /* tp_clear */
2948 0, /* tp_richcompare */
2949 0, /* tp_weaklistoffset */
2950 PyObject_SelfIter, /* tp_iter */
2951 (iternextfunc)count_next, /* tp_iternext */
2952 0, /* tp_methods */
2953 0, /* tp_members */
2954 0, /* tp_getset */
2955 0, /* tp_base */
2956 0, /* tp_dict */
2957 0, /* tp_descr_get */
2958 0, /* tp_descr_set */
2959 0, /* tp_dictoffset */
2960 0, /* tp_init */
2961 0, /* tp_alloc */
2962 count_new, /* tp_new */
2963 };
2964
2965
2966 /* izip object ************************************************************/
2967
2968 #include "Python.h"
2969
2970 typedef struct {
2971 PyObject_HEAD
2972 Py_ssize_t tuplesize;
2973 PyObject *ittuple; /* tuple of iterators */
2974 PyObject *result;
2975 } izipobject;
2976
2977 static PyTypeObject izip_type;
2978
2979 static PyObject *
2980 izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2981 {
2982 izipobject *lz;
2983 Py_ssize_t i;
2984 PyObject *ittuple; /* tuple of iterators */
2985 PyObject *result;
2986 Py_ssize_t tuplesize = PySequence_Length(args);
2987
2988 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
2989 return NULL;
2990
2991 /* args must be a tuple */
2992 assert(PyTuple_Check(args));
2993
2994 /* obtain iterators */
2995 ittuple = PyTuple_New(tuplesize);
2996 if (ittuple == NULL)
2997 return NULL;
2998 for (i=0; i < tuplesize; ++i) {
2999 PyObject *item = PyTuple_GET_ITEM(args, i);
3000 PyObject *it = PyObject_GetIter(item);
3001 if (it == NULL) {
3002 if (PyErr_ExceptionMatches(PyExc_TypeError))
3003 PyErr_Format(PyExc_TypeError,
3004 "izip argument #%zd must support iteration",
3005 i+1);
3006 Py_DECREF(ittuple);
3007 return NULL;
3008 }
3009 PyTuple_SET_ITEM(ittuple, i, it);
3010 }
3011
3012 /* create a result holder */
3013 result = PyTuple_New(tuplesize);
3014 if (result == NULL) {
3015 Py_DECREF(ittuple);
3016 return NULL;
3017 }
3018 for (i=0 ; i < tuplesize ; i++) {
3019 Py_INCREF(Py_None);
3020 PyTuple_SET_ITEM(result, i, Py_None);
3021 }
3022
3023 /* create izipobject structure */
3024 lz = (izipobject *)type->tp_alloc(type, 0);
3025 if (lz == NULL) {
3026 Py_DECREF(ittuple);
3027 Py_DECREF(result);
3028 return NULL;
3029 }
3030 lz->ittuple = ittuple;
3031 lz->tuplesize = tuplesize;
3032 lz->result = result;
3033
3034 return (PyObject *)lz;
3035 }
3036
3037 static void
3038 izip_dealloc(izipobject *lz)
3039 {
3040 PyObject_GC_UnTrack(lz);
3041 Py_XDECREF(lz->ittuple);
3042 Py_XDECREF(lz->result);
3043 Py_TYPE(lz)->tp_free(lz);
3044 }
3045
3046 static int
3047 izip_traverse(izipobject *lz, visitproc visit, void *arg)
3048 {
3049 Py_VISIT(lz->ittuple);
3050 Py_VISIT(lz->result);
3051 return 0;
3052 }
3053
3054 static PyObject *
3055 izip_next(izipobject *lz)
3056 {
3057 Py_ssize_t i;
3058 Py_ssize_t tuplesize = lz->tuplesize;
3059 PyObject *result = lz->result;
3060 PyObject *it;
3061 PyObject *item;
3062 PyObject *olditem;
3063
3064 if (tuplesize == 0)
3065 return NULL;
3066 if (Py_REFCNT(result) == 1) {
3067 Py_INCREF(result);
3068 for (i=0 ; i < tuplesize ; i++) {
3069 it = PyTuple_GET_ITEM(lz->ittuple, i);
3070 assert(PyIter_Check(it));
3071 item = (*Py_TYPE(it)->tp_iternext)(it);
3072 if (item == NULL) {
3073 Py_DECREF(result);
3074 return NULL;
3075 }
3076 olditem = PyTuple_GET_ITEM(result, i);
3077 PyTuple_SET_ITEM(result, i, item);
3078 Py_DECREF(olditem);
3079 }
3080 } else {
3081 result = PyTuple_New(tuplesize);
3082 if (result == NULL)
3083 return NULL;
3084 for (i=0 ; i < tuplesize ; i++) {
3085 it = PyTuple_GET_ITEM(lz->ittuple, i);
3086 assert(PyIter_Check(it));
3087 item = (*Py_TYPE(it)->tp_iternext)(it);
3088 if (item == NULL) {
3089 Py_DECREF(result);
3090 return NULL;
3091 }
3092 PyTuple_SET_ITEM(result, i, item);
3093 }
3094 }
3095 return result;
3096 }
3097
3098 PyDoc_STRVAR(izip_doc,
3099 "izip(iter1 [,iter2 [...]]) --> izip object\n\
3100 \n\
3101 Return a izip object whose .next() method returns a tuple where\n\
3102 the i-th element comes from the i-th iterable argument. The .next()\n\
3103 method continues until the shortest iterable in the argument sequence\n\
3104 is exhausted and then it raises StopIteration. Works like the zip()\n\
3105 function but consumes less memory by returning an iterator instead of\n\
3106 a list.");
3107
3108 static PyTypeObject izip_type = {
3109 PyVarObject_HEAD_INIT(NULL, 0)
3110 "itertools.izip", /* tp_name */
3111 sizeof(izipobject), /* tp_basicsize */
3112 0, /* tp_itemsize */
3113 /* methods */
3114 (destructor)izip_dealloc, /* tp_dealloc */
3115 0, /* tp_print */
3116 0, /* tp_getattr */
3117 0, /* tp_setattr */
3118 0, /* tp_compare */
3119 0, /* tp_repr */
3120 0, /* tp_as_number */
3121 0, /* tp_as_sequence */
3122 0, /* tp_as_mapping */
3123 0, /* tp_hash */
3124 0, /* tp_call */
3125 0, /* tp_str */
3126 PyObject_GenericGetAttr, /* tp_getattro */
3127 0, /* tp_setattro */
3128 0, /* tp_as_buffer */
3129 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3130 Py_TPFLAGS_BASETYPE, /* tp_flags */
3131 izip_doc, /* tp_doc */
3132 (traverseproc)izip_traverse, /* tp_traverse */
3133 0, /* tp_clear */
3134 0, /* tp_richcompare */
3135 0, /* tp_weaklistoffset */
3136 PyObject_SelfIter, /* tp_iter */
3137 (iternextfunc)izip_next, /* tp_iternext */
3138 0, /* tp_methods */
3139 0, /* tp_members */
3140 0, /* tp_getset */
3141 0, /* tp_base */
3142 0, /* tp_dict */
3143 0, /* tp_descr_get */
3144 0, /* tp_descr_set */
3145 0, /* tp_dictoffset */
3146 0, /* tp_init */
3147 0, /* tp_alloc */
3148 izip_new, /* tp_new */
3149 PyObject_GC_Del, /* tp_free */
3150 };
3151
3152
3153 /* repeat object ************************************************************/
3154
3155 typedef struct {
3156 PyObject_HEAD
3157 PyObject *element;
3158 Py_ssize_t cnt;
3159 } repeatobject;
3160
3161 static PyTypeObject repeat_type;
3162
3163 static PyObject *
3164 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3165 {
3166 repeatobject *ro;
3167 PyObject *element;
3168 Py_ssize_t cnt = -1;
3169
3170 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
3171 return NULL;
3172
3173 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
3174 return NULL;
3175
3176 if (PyTuple_Size(args) == 2 && cnt < 0)
3177 cnt = 0;
3178
3179 ro = (repeatobject *)type->tp_alloc(type, 0);
3180 if (ro == NULL)
3181 return NULL;
3182 Py_INCREF(element);
3183 ro->element = element;
3184 ro->cnt = cnt;
3185 return (PyObject *)ro;
3186 }
3187
3188 static void
3189 repeat_dealloc(repeatobject *ro)
3190 {
3191 PyObject_GC_UnTrack(ro);
3192 Py_XDECREF(ro->element);
3193 Py_TYPE(ro)->tp_free(ro);
3194 }
3195
3196 static int
3197 repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3198 {
3199 Py_VISIT(ro->element);
3200 return 0;
3201 }
3202
3203 static PyObject *
3204 repeat_next(repeatobject *ro)
3205 {
3206 if (ro->cnt == 0)
3207 return NULL;
3208 if (ro->cnt > 0)
3209 ro->cnt--;
3210 Py_INCREF(ro->element);
3211 return ro->element;
3212 }
3213
3214 static PyObject *
3215 repeat_repr(repeatobject *ro)
3216 {
3217 PyObject *result, *objrepr;
3218
3219 objrepr = PyObject_Repr(ro->element);
3220 if (objrepr == NULL)
3221 return NULL;
3222
3223 if (ro->cnt == -1)
3224 result = PyString_FromFormat("repeat(%s)",
3225 PyString_AS_STRING(objrepr));
3226 else
3227 result = PyString_FromFormat("repeat(%s, %zd)",
3228 PyString_AS_STRING(objrepr), ro->cnt);
3229 Py_DECREF(objrepr);
3230 return result;
3231 }
3232
3233 static PyObject *
3234 repeat_len(repeatobject *ro)
3235 {
3236 if (ro->cnt == -1) {
3237 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3238 return NULL;
3239 }
3240 return PyInt_FromSize_t(ro->cnt);
3241 }
3242
3243 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
3244
3245 static PyMethodDef repeat_methods[] = {
3246 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3247 {NULL, NULL} /* sentinel */
3248 };
3249
3250 PyDoc_STRVAR(repeat_doc,
3251 "repeat(element [,times]) -> create an iterator which returns the element\n\
3252 for the specified number of times. If not specified, returns the element\n\
3253 endlessly.");
3254
3255 static PyTypeObject repeat_type = {
3256 PyVarObject_HEAD_INIT(NULL, 0)
3257 "itertools.repeat", /* tp_name */
3258 sizeof(repeatobject), /* tp_basicsize */
3259 0, /* tp_itemsize */
3260 /* methods */
3261 (destructor)repeat_dealloc, /* tp_dealloc */
3262 0, /* tp_print */
3263 0, /* tp_getattr */
3264 0, /* tp_setattr */
3265 0, /* tp_compare */
3266 (reprfunc)repeat_repr, /* tp_repr */
3267 0, /* tp_as_number */
3268 0, /* tp_as_sequence */
3269 0, /* tp_as_mapping */
3270 0, /* tp_hash */
3271 0, /* tp_call */
3272 0, /* tp_str */
3273 PyObject_GenericGetAttr, /* tp_getattro */
3274 0, /* tp_setattro */
3275 0, /* tp_as_buffer */
3276 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3277 Py_TPFLAGS_BASETYPE, /* tp_flags */
3278 repeat_doc, /* tp_doc */
3279 (traverseproc)repeat_traverse, /* tp_traverse */
3280 0, /* tp_clear */
3281 0, /* tp_richcompare */
3282 0, /* tp_weaklistoffset */
3283 PyObject_SelfIter, /* tp_iter */
3284 (iternextfunc)repeat_next, /* tp_iternext */
3285 repeat_methods, /* tp_methods */
3286 0, /* tp_members */
3287 0, /* tp_getset */
3288 0, /* tp_base */
3289 0, /* tp_dict */
3290 0, /* tp_descr_get */
3291 0, /* tp_descr_set */
3292 0, /* tp_dictoffset */
3293 0, /* tp_init */
3294 0, /* tp_alloc */
3295 repeat_new, /* tp_new */
3296 PyObject_GC_Del, /* tp_free */
3297 };
3298
3299 /* iziplongest object ************************************************************/
3300
3301 #include "Python.h"
3302
3303 typedef struct {
3304 PyObject_HEAD
3305 Py_ssize_t tuplesize;
3306 Py_ssize_t numactive;
3307 PyObject *ittuple; /* tuple of iterators */
3308 PyObject *result;
3309 PyObject *fillvalue;
3310 } iziplongestobject;
3311
3312 static PyTypeObject iziplongest_type;
3313
3314 static PyObject *
3315 izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3316 {
3317 iziplongestobject *lz;
3318 Py_ssize_t i;
3319 PyObject *ittuple; /* tuple of iterators */
3320 PyObject *result;
3321 PyObject *fillvalue = Py_None;
3322 Py_ssize_t tuplesize = PySequence_Length(args);
3323
3324 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3325 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3326 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3327 PyErr_SetString(PyExc_TypeError,
3328 "izip_longest() got an unexpected keyword argument");
3329 return NULL;
3330 }
3331 }
3332
3333 /* args must be a tuple */
3334 assert(PyTuple_Check(args));
3335
3336 /* obtain iterators */
3337 ittuple = PyTuple_New(tuplesize);
3338 if (ittuple == NULL)
3339 return NULL;
3340 for (i=0; i < tuplesize; ++i) {
3341 PyObject *item = PyTuple_GET_ITEM(args, i);
3342 PyObject *it = PyObject_GetIter(item);
3343 if (it == NULL) {
3344 if (PyErr_ExceptionMatches(PyExc_TypeError))
3345 PyErr_Format(PyExc_TypeError,
3346 "izip_longest argument #%zd must support iteration",
3347 i+1);
3348 Py_DECREF(ittuple);
3349 return NULL;
3350 }
3351 PyTuple_SET_ITEM(ittuple, i, it);
3352 }
3353
3354 /* create a result holder */
3355 result = PyTuple_New(tuplesize);
3356 if (result == NULL) {
3357 Py_DECREF(ittuple);
3358 return NULL;
3359 }
3360 for (i=0 ; i < tuplesize ; i++) {
3361 Py_INCREF(Py_None);
3362 PyTuple_SET_ITEM(result, i, Py_None);
3363 }
3364
3365 /* create iziplongestobject structure */
3366 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3367 if (lz == NULL) {
3368 Py_DECREF(ittuple);
3369 Py_DECREF(result);
3370 return NULL;
3371 }
3372 lz->ittuple = ittuple;
3373 lz->tuplesize = tuplesize;
3374 lz->numactive = tuplesize;
3375 lz->result = result;
3376 Py_INCREF(fillvalue);
3377 lz->fillvalue = fillvalue;
3378 return (PyObject *)lz;
3379 }
3380
3381 static void
3382 izip_longest_dealloc(iziplongestobject *lz)
3383 {
3384 PyObject_GC_UnTrack(lz);
3385 Py_XDECREF(lz->ittuple);
3386 Py_XDECREF(lz->result);
3387 Py_XDECREF(lz->fillvalue);
3388 Py_TYPE(lz)->tp_free(lz);
3389 }
3390
3391 static int
3392 izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3393 {
3394 Py_VISIT(lz->ittuple);
3395 Py_VISIT(lz->result);
3396 Py_VISIT(lz->fillvalue);
3397 return 0;
3398 }
3399
3400 static PyObject *
3401 izip_longest_next(iziplongestobject *lz)
3402 {
3403 Py_ssize_t i;
3404 Py_ssize_t tuplesize = lz->tuplesize;
3405 PyObject *result = lz->result;
3406 PyObject *it;
3407 PyObject *item;
3408 PyObject *olditem;
3409
3410 if (tuplesize == 0)
3411 return NULL;
3412 if (lz->numactive == 0)
3413 return NULL;
3414 if (Py_REFCNT(result) == 1) {
3415 Py_INCREF(result);
3416 for (i=0 ; i < tuplesize ; i++) {
3417 it = PyTuple_GET_ITEM(lz->ittuple, i);
3418 if (it == NULL) {
3419 Py_INCREF(lz->fillvalue);
3420 item = lz->fillvalue;
3421 } else {
3422 assert(PyIter_Check(it));
3423 item = (*Py_TYPE(it)->tp_iternext)(it);
3424 if (item == NULL) {
3425 lz->numactive -= 1;
3426 if (lz->numactive == 0) {
3427 Py_DECREF(result);
3428 return NULL;
3429 } else {
3430 Py_INCREF(lz->fillvalue);
3431 item = lz->fillvalue;
3432 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3433 Py_DECREF(it);
3434 }
3435 }
3436 }
3437 olditem = PyTuple_GET_ITEM(result, i);
3438 PyTuple_SET_ITEM(result, i, item);
3439 Py_DECREF(olditem);
3440 }
3441 } else {
3442 result = PyTuple_New(tuplesize);
3443 if (result == NULL)
3444 return NULL;
3445 for (i=0 ; i < tuplesize ; i++) {
3446 it = PyTuple_GET_ITEM(lz->ittuple, i);
3447 if (it == NULL) {
3448 Py_INCREF(lz->fillvalue);
3449 item = lz->fillvalue;
3450 } else {
3451 assert(PyIter_Check(it));
3452 item = (*Py_TYPE(it)->tp_iternext)(it);
3453 if (item == NULL) {
3454 lz->numactive -= 1;
3455 if (lz->numactive == 0) {
3456 Py_DECREF(result);
3457 return NULL;
3458 } else {
3459 Py_INCREF(lz->fillvalue);
3460 item = lz->fillvalue;
3461 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3462 Py_DECREF(it);
3463 }
3464 }
3465 }
3466 PyTuple_SET_ITEM(result, i, item);
3467 }
3468 }
3469 return result;
3470 }
3471
3472 PyDoc_STRVAR(izip_longest_doc,
3473 "izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3474 \n\
3475 Return an izip_longest object whose .next() method returns a tuple where\n\
3476 the i-th element comes from the i-th iterable argument. The .next()\n\
3477 method continues until the longest iterable in the argument sequence\n\
3478 is exhausted and then it raises StopIteration. When the shorter iterables\n\
3479 are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3480 defaults to None or can be specified by a keyword argument.\n\
3481 ");
3482
3483 static PyTypeObject iziplongest_type = {
3484 PyVarObject_HEAD_INIT(NULL, 0)
3485 "itertools.izip_longest", /* tp_name */
3486 sizeof(iziplongestobject), /* tp_basicsize */
3487 0, /* tp_itemsize */
3488 /* methods */
3489 (destructor)izip_longest_dealloc, /* tp_dealloc */
3490 0, /* tp_print */
3491 0, /* tp_getattr */
3492 0, /* tp_setattr */
3493 0, /* tp_compare */
3494 0, /* tp_repr */
3495 0, /* tp_as_number */
3496 0, /* tp_as_sequence */
3497 0, /* tp_as_mapping */
3498 0, /* tp_hash */
3499 0, /* tp_call */
3500 0, /* tp_str */
3501 PyObject_GenericGetAttr, /* tp_getattro */
3502 0, /* tp_setattro */
3503 0, /* tp_as_buffer */
3504 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3505 Py_TPFLAGS_BASETYPE, /* tp_flags */
3506 izip_longest_doc, /* tp_doc */
3507 (traverseproc)izip_longest_traverse, /* tp_traverse */
3508 0, /* tp_clear */
3509 0, /* tp_richcompare */
3510 0, /* tp_weaklistoffset */
3511 PyObject_SelfIter, /* tp_iter */
3512 (iternextfunc)izip_longest_next, /* tp_iternext */
3513 0, /* tp_methods */
3514 0, /* tp_members */
3515 0, /* tp_getset */
3516 0, /* tp_base */
3517 0, /* tp_dict */
3518 0, /* tp_descr_get */
3519 0, /* tp_descr_set */
3520 0, /* tp_dictoffset */
3521 0, /* tp_init */
3522 0, /* tp_alloc */
3523 izip_longest_new, /* tp_new */
3524 PyObject_GC_Del, /* tp_free */
3525 };
3526
3527 /* module level code ********************************************************/
3528
3529 PyDoc_STRVAR(module_doc,
3530 "Functional tools for creating and using iterators.\n\
3531 \n\
3532 Infinite iterators:\n\
3533 count([n]) --> n, n+1, n+2, ...\n\
3534 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
3535 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
3536 \n\
3537 Iterators terminating on the shortest input sequence:\n\
3538 izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3539 izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
3540 ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
3541 ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
3542 islice(seq, [start,] stop [, step]) --> elements from\n\
3543 seq[start:stop:step]\n\
3544 imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
3545 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
3546 tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
3547 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
3548 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
3549 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
3550 groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
3551 ");
3552
3553
3554 static PyMethodDef module_methods[] = {
3555 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
3556 {NULL, NULL} /* sentinel */
3557 };
3558
3559 PyMODINIT_FUNC
3560 inititertools(void)
3561 {
3562 int i;
3563 PyObject *m;
3564 char *name;
3565 PyTypeObject *typelist[] = {
3566 &combinations_type,
3567 &cycle_type,
3568 &dropwhile_type,
3569 &takewhile_type,
3570 &islice_type,
3571 &starmap_type,
3572 &imap_type,
3573 &chain_type,
3574 &ifilter_type,
3575 &ifilterfalse_type,
3576 &count_type,
3577 &izip_type,
3578 &iziplongest_type,
3579 &permutations_type,
3580 &product_type,
3581 &repeat_type,
3582 &groupby_type,
3583 NULL
3584 };
3585
3586 Py_TYPE(&teedataobject_type) = &PyType_Type;
3587 m = Py_InitModule3("itertools", module_methods, module_doc);
3588 if (m == NULL)
3589 return;
3590
3591 for (i=0 ; typelist[i] != NULL ; i++) {
3592 if (PyType_Ready(typelist[i]) < 0)
3593 return;
3594 name = strchr(typelist[i]->tp_name, '.');
3595 assert (name != NULL);
3596 Py_INCREF(typelist[i]);
3597 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
3598 }
3599
3600 if (PyType_Ready(&teedataobject_type) < 0)
3601 return;
3602 if (PyType_Ready(&tee_type) < 0)
3603 return;
3604 if (PyType_Ready(&_grouper_type) < 0)
3605 return;
3606 }
python git mirror