?  search:
summary | log | graphiclog1 | graphiclog2 | commit | commitdiff | tree | refs | edit | fork
blame | history | raw | HEAD
Issue 6305: Clarify error message for large arguments to itertools.islice().
[python.git] / Modules / itertoolsmodule.c
blob64e36066b4780c2444b53222f858a0326b4514a4
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 && PyList_Append(lz->saved, item)) {
749 Py_DECREF(item);
750 return NULL;
751 }
752 return item;
753 }
754 if (PyErr_Occurred()) {
755 if (PyErr_ExceptionMatches(PyExc_StopIteration))
756 PyErr_Clear();
757 else
758 return NULL;
759 }
760 if (PyList_Size(lz->saved) == 0)
761 return NULL;
762 it = PyObject_GetIter(lz->saved);
763 if (it == NULL)
764 return NULL;
765 tmp = lz->it;
766 lz->it = it;
767 lz->firstpass = 1;
768 Py_DECREF(tmp);
769 }
770 }
771
772 PyDoc_STRVAR(cycle_doc,
773 "cycle(iterable) --> cycle object\n\
774 \n\
775 Return elements from the iterable until it is exhausted.\n\
776 Then repeat the sequence indefinitely.");
777
778 static PyTypeObject cycle_type = {
779 PyVarObject_HEAD_INIT(NULL, 0)
780 "itertools.cycle", /* tp_name */
781 sizeof(cycleobject), /* tp_basicsize */
782 0, /* tp_itemsize */
783 /* methods */
784 (destructor)cycle_dealloc, /* tp_dealloc */
785 0, /* tp_print */
786 0, /* tp_getattr */
787 0, /* tp_setattr */
788 0, /* tp_compare */
789 0, /* tp_repr */
790 0, /* tp_as_number */
791 0, /* tp_as_sequence */
792 0, /* tp_as_mapping */
793 0, /* tp_hash */
794 0, /* tp_call */
795 0, /* tp_str */
796 PyObject_GenericGetAttr, /* tp_getattro */
797 0, /* tp_setattro */
798 0, /* tp_as_buffer */
799 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
800 Py_TPFLAGS_BASETYPE, /* tp_flags */
801 cycle_doc, /* tp_doc */
802 (traverseproc)cycle_traverse, /* tp_traverse */
803 0, /* tp_clear */
804 0, /* tp_richcompare */
805 0, /* tp_weaklistoffset */
806 PyObject_SelfIter, /* tp_iter */
807 (iternextfunc)cycle_next, /* tp_iternext */
808 0, /* tp_methods */
809 0, /* tp_members */
810 0, /* tp_getset */
811 0, /* tp_base */
812 0, /* tp_dict */
813 0, /* tp_descr_get */
814 0, /* tp_descr_set */
815 0, /* tp_dictoffset */
816 0, /* tp_init */
817 0, /* tp_alloc */
818 cycle_new, /* tp_new */
819 PyObject_GC_Del, /* tp_free */
820 };
821
822
823 /* dropwhile object **********************************************************/
824
825 typedef struct {
826 PyObject_HEAD
827 PyObject *func;
828 PyObject *it;
829 long start;
830 } dropwhileobject;
831
832 static PyTypeObject dropwhile_type;
833
834 static PyObject *
835 dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
836 {
837 PyObject *func, *seq;
838 PyObject *it;
839 dropwhileobject *lz;
840
841 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
842 return NULL;
843
844 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
845 return NULL;
846
847 /* Get iterator. */
848 it = PyObject_GetIter(seq);
849 if (it == NULL)
850 return NULL;
851
852 /* create dropwhileobject structure */
853 lz = (dropwhileobject *)type->tp_alloc(type, 0);
854 if (lz == NULL) {
855 Py_DECREF(it);
856 return NULL;
857 }
858 Py_INCREF(func);
859 lz->func = func;
860 lz->it = it;
861 lz->start = 0;
862
863 return (PyObject *)lz;
864 }
865
866 static void
867 dropwhile_dealloc(dropwhileobject *lz)
868 {
869 PyObject_GC_UnTrack(lz);
870 Py_XDECREF(lz->func);
871 Py_XDECREF(lz->it);
872 Py_TYPE(lz)->tp_free(lz);
873 }
874
875 static int
876 dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
877 {
878 Py_VISIT(lz->it);
879 Py_VISIT(lz->func);
880 return 0;
881 }
882
883 static PyObject *
884 dropwhile_next(dropwhileobject *lz)
885 {
886 PyObject *item, *good;
887 PyObject *it = lz->it;
888 long ok;
889 PyObject *(*iternext)(PyObject *);
890
891 iternext = *Py_TYPE(it)->tp_iternext;
892 for (;;) {
893 item = iternext(it);
894 if (item == NULL)
895 return NULL;
896 if (lz->start == 1)
897 return item;
898
899 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
900 if (good == NULL) {
901 Py_DECREF(item);
902 return NULL;
903 }
904 ok = PyObject_IsTrue(good);
905 Py_DECREF(good);
906 if (!ok) {
907 lz->start = 1;
908 return item;
909 }
910 Py_DECREF(item);
911 }
912 }
913
914 PyDoc_STRVAR(dropwhile_doc,
915 "dropwhile(predicate, iterable) --> dropwhile object\n\
916 \n\
917 Drop items from the iterable while predicate(item) is true.\n\
918 Afterwards, return every element until the iterable is exhausted.");
919
920 static PyTypeObject dropwhile_type = {
921 PyVarObject_HEAD_INIT(NULL, 0)
922 "itertools.dropwhile", /* tp_name */
923 sizeof(dropwhileobject), /* tp_basicsize */
924 0, /* tp_itemsize */
925 /* methods */
926 (destructor)dropwhile_dealloc, /* tp_dealloc */
927 0, /* tp_print */
928 0, /* tp_getattr */
929 0, /* tp_setattr */
930 0, /* tp_compare */
931 0, /* tp_repr */
932 0, /* tp_as_number */
933 0, /* tp_as_sequence */
934 0, /* tp_as_mapping */
935 0, /* tp_hash */
936 0, /* tp_call */
937 0, /* tp_str */
938 PyObject_GenericGetAttr, /* tp_getattro */
939 0, /* tp_setattro */
940 0, /* tp_as_buffer */
941 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
942 Py_TPFLAGS_BASETYPE, /* tp_flags */
943 dropwhile_doc, /* tp_doc */
944 (traverseproc)dropwhile_traverse, /* tp_traverse */
945 0, /* tp_clear */
946 0, /* tp_richcompare */
947 0, /* tp_weaklistoffset */
948 PyObject_SelfIter, /* tp_iter */
949 (iternextfunc)dropwhile_next, /* tp_iternext */
950 0, /* tp_methods */
951 0, /* tp_members */
952 0, /* tp_getset */
953 0, /* tp_base */
954 0, /* tp_dict */
955 0, /* tp_descr_get */
956 0, /* tp_descr_set */
957 0, /* tp_dictoffset */
958 0, /* tp_init */
959 0, /* tp_alloc */
960 dropwhile_new, /* tp_new */
961 PyObject_GC_Del, /* tp_free */
962 };
963
964
965 /* takewhile object **********************************************************/
966
967 typedef struct {
968 PyObject_HEAD
969 PyObject *func;
970 PyObject *it;
971 long stop;
972 } takewhileobject;
973
974 static PyTypeObject takewhile_type;
975
976 static PyObject *
977 takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
978 {
979 PyObject *func, *seq;
980 PyObject *it;
981 takewhileobject *lz;
982
983 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
984 return NULL;
985
986 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
987 return NULL;
988
989 /* Get iterator. */
990 it = PyObject_GetIter(seq);
991 if (it == NULL)
992 return NULL;
993
994 /* create takewhileobject structure */
995 lz = (takewhileobject *)type->tp_alloc(type, 0);
996 if (lz == NULL) {
997 Py_DECREF(it);
998 return NULL;
999 }
1000 Py_INCREF(func);
1001 lz->func = func;
1002 lz->it = it;
1003 lz->stop = 0;
1004
1005 return (PyObject *)lz;
1006 }
1007
1008 static void
1009 takewhile_dealloc(takewhileobject *lz)
1010 {
1011 PyObject_GC_UnTrack(lz);
1012 Py_XDECREF(lz->func);
1013 Py_XDECREF(lz->it);
1014 Py_TYPE(lz)->tp_free(lz);
1015 }
1016
1017 static int
1018 takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1019 {
1020 Py_VISIT(lz->it);
1021 Py_VISIT(lz->func);
1022 return 0;
1023 }
1024
1025 static PyObject *
1026 takewhile_next(takewhileobject *lz)
1027 {
1028 PyObject *item, *good;
1029 PyObject *it = lz->it;
1030 long ok;
1031
1032 if (lz->stop == 1)
1033 return NULL;
1034
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 None or an integer: 0 <= x <= maxint.");
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 None or an integer: 0 <= x <= maxint.");
1156 return NULL;
1157 }
1158 }
1159 }
1160 if (start<0 || stop<-1) {
1161 PyErr_SetString(PyExc_ValueError,
1162 "Indices for islice() must be None or an integer: 0 <= x <= maxint.");
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 iternext = *Py_TYPE(it)->tp_iternext;
1222 while (lz->cnt < lz->next) {
1223 item = iternext(it);
1224 if (item == NULL)
1225 return NULL;
1226 Py_DECREF(item);
1227 lz->cnt++;
1228 }
1229 if (lz->stop != -1 && lz->cnt >= lz->stop)
1230 return NULL;
1231 item = iternext(it);
1232 if (item == NULL)
1233 return NULL;
1234 lz->cnt++;
1235 oldnext = lz->next;
1236 lz->next += lz->step;
1237 if (lz->next < oldnext) /* Check for overflow */
1238 lz->next = lz->stop;
1239 return item;
1240 }
1241
1242 PyDoc_STRVAR(islice_doc,
1243 "islice(iterable, [start,] stop [, step]) --> islice object\n\
1244 \n\
1245 Return an iterator whose next() method returns selected values from an\n\
1246 iterable. If start is specified, will skip all preceding elements;\n\
1247 otherwise, start defaults to zero. Step defaults to one. If\n\
1248 specified as another value, step determines how many values are \n\
1249 skipped between successive calls. Works like a slice() on a list\n\
1250 but returns an iterator.");
1251
1252 static PyTypeObject islice_type = {
1253 PyVarObject_HEAD_INIT(NULL, 0)
1254 "itertools.islice", /* tp_name */
1255 sizeof(isliceobject), /* tp_basicsize */
1256 0, /* tp_itemsize */
1257 /* methods */
1258 (destructor)islice_dealloc, /* tp_dealloc */
1259 0, /* tp_print */
1260 0, /* tp_getattr */
1261 0, /* tp_setattr */
1262 0, /* tp_compare */
1263 0, /* tp_repr */
1264 0, /* tp_as_number */
1265 0, /* tp_as_sequence */
1266 0, /* tp_as_mapping */
1267 0, /* tp_hash */
1268 0, /* tp_call */
1269 0, /* tp_str */
1270 PyObject_GenericGetAttr, /* tp_getattro */
1271 0, /* tp_setattro */
1272 0, /* tp_as_buffer */
1273 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1274 Py_TPFLAGS_BASETYPE, /* tp_flags */
1275 islice_doc, /* tp_doc */
1276 (traverseproc)islice_traverse, /* tp_traverse */
1277 0, /* tp_clear */
1278 0, /* tp_richcompare */
1279 0, /* tp_weaklistoffset */
1280 PyObject_SelfIter, /* tp_iter */
1281 (iternextfunc)islice_next, /* tp_iternext */
1282 0, /* tp_methods */
1283 0, /* tp_members */
1284 0, /* tp_getset */
1285 0, /* tp_base */
1286 0, /* tp_dict */
1287 0, /* tp_descr_get */
1288 0, /* tp_descr_set */
1289 0, /* tp_dictoffset */
1290 0, /* tp_init */
1291 0, /* tp_alloc */
1292 islice_new, /* tp_new */
1293 PyObject_GC_Del, /* tp_free */
1294 };
1295
1296
1297 /* starmap object ************************************************************/
1298
1299 typedef struct {
1300 PyObject_HEAD
1301 PyObject *func;
1302 PyObject *it;
1303 } starmapobject;
1304
1305 static PyTypeObject starmap_type;
1306
1307 static PyObject *
1308 starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1309 {
1310 PyObject *func, *seq;
1311 PyObject *it;
1312 starmapobject *lz;
1313
1314 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
1315 return NULL;
1316
1317 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1318 return NULL;
1319
1320 /* Get iterator. */
1321 it = PyObject_GetIter(seq);
1322 if (it == NULL)
1323 return NULL;
1324
1325 /* create starmapobject structure */
1326 lz = (starmapobject *)type->tp_alloc(type, 0);
1327 if (lz == NULL) {
1328 Py_DECREF(it);
1329 return NULL;
1330 }
1331 Py_INCREF(func);
1332 lz->func = func;
1333 lz->it = it;
1334
1335 return (PyObject *)lz;
1336 }
1337
1338 static void
1339 starmap_dealloc(starmapobject *lz)
1340 {
1341 PyObject_GC_UnTrack(lz);
1342 Py_XDECREF(lz->func);
1343 Py_XDECREF(lz->it);
1344 Py_TYPE(lz)->tp_free(lz);
1345 }
1346
1347 static int
1348 starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1349 {
1350 Py_VISIT(lz->it);
1351 Py_VISIT(lz->func);
1352 return 0;
1353 }
1354
1355 static PyObject *
1356 starmap_next(starmapobject *lz)
1357 {
1358 PyObject *args;
1359 PyObject *result;
1360 PyObject *it = lz->it;
1361
1362 args = (*Py_TYPE(it)->tp_iternext)(it);
1363 if (args == NULL)
1364 return NULL;
1365 if (!PyTuple_CheckExact(args)) {
1366 PyObject *newargs = PySequence_Tuple(args);
1367 Py_DECREF(args);
1368 if (newargs == NULL)
1369 return NULL;
1370 args = newargs;
1371 }
1372 result = PyObject_Call(lz->func, args, NULL);
1373 Py_DECREF(args);
1374 return result;
1375 }
1376
1377 PyDoc_STRVAR(starmap_doc,
1378 "starmap(function, sequence) --> starmap object\n\
1379 \n\
1380 Return an iterator whose values are returned from the function evaluated\n\
1381 with a argument tuple taken from the given sequence.");
1382
1383 static PyTypeObject starmap_type = {
1384 PyVarObject_HEAD_INIT(NULL, 0)
1385 "itertools.starmap", /* tp_name */
1386 sizeof(starmapobject), /* tp_basicsize */
1387 0, /* tp_itemsize */
1388 /* methods */
1389 (destructor)starmap_dealloc, /* tp_dealloc */
1390 0, /* tp_print */
1391 0, /* tp_getattr */
1392 0, /* tp_setattr */
1393 0, /* tp_compare */
1394 0, /* tp_repr */
1395 0, /* tp_as_number */
1396 0, /* tp_as_sequence */
1397 0, /* tp_as_mapping */
1398 0, /* tp_hash */
1399 0, /* tp_call */
1400 0, /* tp_str */
1401 PyObject_GenericGetAttr, /* tp_getattro */
1402 0, /* tp_setattro */
1403 0, /* tp_as_buffer */
1404 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1405 Py_TPFLAGS_BASETYPE, /* tp_flags */
1406 starmap_doc, /* tp_doc */
1407 (traverseproc)starmap_traverse, /* tp_traverse */
1408 0, /* tp_clear */
1409 0, /* tp_richcompare */
1410 0, /* tp_weaklistoffset */
1411 PyObject_SelfIter, /* tp_iter */
1412 (iternextfunc)starmap_next, /* tp_iternext */
1413 0, /* tp_methods */
1414 0, /* tp_members */
1415 0, /* tp_getset */
1416 0, /* tp_base */
1417 0, /* tp_dict */
1418 0, /* tp_descr_get */
1419 0, /* tp_descr_set */
1420 0, /* tp_dictoffset */
1421 0, /* tp_init */
1422 0, /* tp_alloc */
1423 starmap_new, /* tp_new */
1424 PyObject_GC_Del, /* tp_free */
1425 };
1426
1427
1428 /* imap object ************************************************************/
1429
1430 typedef struct {
1431 PyObject_HEAD
1432 PyObject *iters;
1433 PyObject *func;
1434 } imapobject;
1435
1436 static PyTypeObject imap_type;
1437
1438 static PyObject *
1439 imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1440 {
1441 PyObject *it, *iters, *func;
1442 imapobject *lz;
1443 Py_ssize_t numargs, i;
1444
1445 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
1446 return NULL;
1447
1448 numargs = PyTuple_Size(args);
1449 if (numargs < 2) {
1450 PyErr_SetString(PyExc_TypeError,
1451 "imap() must have at least two arguments.");
1452 return NULL;
1453 }
1454
1455 iters = PyTuple_New(numargs-1);
1456 if (iters == NULL)
1457 return NULL;
1458
1459 for (i=1 ; i<numargs ; i++) {
1460 /* Get iterator. */
1461 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1462 if (it == NULL) {
1463 Py_DECREF(iters);
1464 return NULL;
1465 }
1466 PyTuple_SET_ITEM(iters, i-1, it);
1467 }
1468
1469 /* create imapobject structure */
1470 lz = (imapobject *)type->tp_alloc(type, 0);
1471 if (lz == NULL) {
1472 Py_DECREF(iters);
1473 return NULL;
1474 }
1475 lz->iters = iters;
1476 func = PyTuple_GET_ITEM(args, 0);
1477 Py_INCREF(func);
1478 lz->func = func;
1479
1480 return (PyObject *)lz;
1481 }
1482
1483 static void
1484 imap_dealloc(imapobject *lz)
1485 {
1486 PyObject_GC_UnTrack(lz);
1487 Py_XDECREF(lz->iters);
1488 Py_XDECREF(lz->func);
1489 Py_TYPE(lz)->tp_free(lz);
1490 }
1491
1492 static int
1493 imap_traverse(imapobject *lz, visitproc visit, void *arg)
1494 {
1495 Py_VISIT(lz->iters);
1496 Py_VISIT(lz->func);
1497 return 0;
1498 }
1499
1500 /*
1501 imap() is an iterator version of __builtins__.map() except that it does
1502 not have the None fill-in feature. That was intentionally left out for
1503 the following reasons:
1504
1505 1) Itertools are designed to be easily combined and chained together.
1506 Having all tools stop with the shortest input is a unifying principle
1507 that makes it easier to combine finite iterators (supplying data) with
1508 infinite iterators like count() and repeat() (for supplying sequential
1509 or constant arguments to a function).
1510
1511 2) In typical use cases for combining itertools, having one finite data
1512 supplier run out before another is likely to be an error condition which
1513 should not pass silently by automatically supplying None.
1514
1515 3) The use cases for automatic None fill-in are rare -- not many functions
1516 do something useful when a parameter suddenly switches type and becomes
1517 None.
1518
1519 4) If a need does arise, it can be met by __builtins__.map() or by
1520 writing: chain(iterable, repeat(None)).
1521
1522 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1523 */
1524
1525 static PyObject *
1526 imap_next(imapobject *lz)
1527 {
1528 PyObject *val;
1529 PyObject *argtuple;
1530 PyObject *result;
1531 Py_ssize_t numargs, i;
1532
1533 numargs = PyTuple_Size(lz->iters);
1534 argtuple = PyTuple_New(numargs);
1535 if (argtuple == NULL)
1536 return NULL;
1537
1538 for (i=0 ; i<numargs ; i++) {
1539 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1540 if (val == NULL) {
1541 Py_DECREF(argtuple);
1542 return NULL;
1543 }
1544 PyTuple_SET_ITEM(argtuple, i, val);
1545 }
1546 if (lz->func == Py_None)
1547 return argtuple;
1548 result = PyObject_Call(lz->func, argtuple, NULL);
1549 Py_DECREF(argtuple);
1550 return result;
1551 }
1552
1553 PyDoc_STRVAR(imap_doc,
1554 "imap(func, *iterables) --> imap object\n\
1555 \n\
1556 Make an iterator that computes the function using arguments from\n\
1557 each of the iterables. Like map() except that it returns\n\
1558 an iterator instead of a list and that it stops when the shortest\n\
1559 iterable is exhausted instead of filling in None for shorter\n\
1560 iterables.");
1561
1562 static PyTypeObject imap_type = {
1563 PyVarObject_HEAD_INIT(NULL, 0)
1564 "itertools.imap", /* tp_name */
1565 sizeof(imapobject), /* tp_basicsize */
1566 0, /* tp_itemsize */
1567 /* methods */
1568 (destructor)imap_dealloc, /* tp_dealloc */
1569 0, /* tp_print */
1570 0, /* tp_getattr */
1571 0, /* tp_setattr */
1572 0, /* tp_compare */
1573 0, /* tp_repr */
1574 0, /* tp_as_number */
1575 0, /* tp_as_sequence */
1576 0, /* tp_as_mapping */
1577 0, /* tp_hash */
1578 0, /* tp_call */
1579 0, /* tp_str */
1580 PyObject_GenericGetAttr, /* tp_getattro */
1581 0, /* tp_setattro */
1582 0, /* tp_as_buffer */
1583 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1584 Py_TPFLAGS_BASETYPE, /* tp_flags */
1585 imap_doc, /* tp_doc */
1586 (traverseproc)imap_traverse, /* tp_traverse */
1587 0, /* tp_clear */
1588 0, /* tp_richcompare */
1589 0, /* tp_weaklistoffset */
1590 PyObject_SelfIter, /* tp_iter */
1591 (iternextfunc)imap_next, /* tp_iternext */
1592 0, /* tp_methods */
1593 0, /* tp_members */
1594 0, /* tp_getset */
1595 0, /* tp_base */
1596 0, /* tp_dict */
1597 0, /* tp_descr_get */
1598 0, /* tp_descr_set */
1599 0, /* tp_dictoffset */
1600 0, /* tp_init */
1601 0, /* tp_alloc */
1602 imap_new, /* tp_new */
1603 PyObject_GC_Del, /* tp_free */
1604 };
1605
1606
1607 /* chain object ************************************************************/
1608
1609 typedef struct {
1610 PyObject_HEAD
1611 PyObject *source; /* Iterator over input iterables */
1612 PyObject *active; /* Currently running input iterator */
1613 } chainobject;
1614
1615 static PyTypeObject chain_type;
1616
1617 static PyObject *
1618 chain_new_internal(PyTypeObject *type, PyObject *source)
1619 {
1620 chainobject *lz;
1621
1622 lz = (chainobject *)type->tp_alloc(type, 0);
1623 if (lz == NULL) {
1624 Py_DECREF(source);
1625 return NULL;
1626 }
1627
1628 lz->source = source;
1629 lz->active = NULL;
1630 return (PyObject *)lz;
1631 }
1632
1633 static PyObject *
1634 chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1635 {
1636 PyObject *source;
1637
1638 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
1639 return NULL;
1640
1641 source = PyObject_GetIter(args);
1642 if (source == NULL)
1643 return NULL;
1644
1645 return chain_new_internal(type, source);
1646 }
1647
1648 static PyObject *
1649 chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
1650 {
1651 PyObject *source;
1652
1653 source = PyObject_GetIter(arg);
1654 if (source == NULL)
1655 return NULL;
1656
1657 return chain_new_internal(type, source);
1658 }
1659
1660 static void
1661 chain_dealloc(chainobject *lz)
1662 {
1663 PyObject_GC_UnTrack(lz);
1664 Py_XDECREF(lz->active);
1665 Py_XDECREF(lz->source);
1666 Py_TYPE(lz)->tp_free(lz);
1667 }
1668
1669 static int
1670 chain_traverse(chainobject *lz, visitproc visit, void *arg)
1671 {
1672 Py_VISIT(lz->source);
1673 Py_VISIT(lz->active);
1674 return 0;
1675 }
1676
1677 static PyObject *
1678 chain_next(chainobject *lz)
1679 {
1680 PyObject *item;
1681
1682 if (lz->source == NULL)
1683 return NULL; /* already stopped */
1684
1685 if (lz->active == NULL) {
1686 PyObject *iterable = PyIter_Next(lz->source);
1687 if (iterable == NULL) {
1688 Py_CLEAR(lz->source);
1689 return NULL; /* no more input sources */
1690 }
1691 lz->active = PyObject_GetIter(iterable);
1692 Py_DECREF(iterable);
1693 if (lz->active == NULL) {
1694 Py_CLEAR(lz->source);
1695 return NULL; /* input not iterable */
1696 }
1697 }
1698 item = PyIter_Next(lz->active);
1699 if (item != NULL)
1700 return item;
1701 if (PyErr_Occurred()) {
1702 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1703 PyErr_Clear();
1704 else
1705 return NULL; /* input raised an exception */
1706 }
1707 Py_CLEAR(lz->active);
1708 return chain_next(lz); /* recurse and use next active */
1709 }
1710
1711 PyDoc_STRVAR(chain_doc,
1712 "chain(*iterables) --> chain object\n\
1713 \n\
1714 Return a chain object whose .next() method returns elements from the\n\
1715 first iterable until it is exhausted, then elements from the next\n\
1716 iterable, until all of the iterables are exhausted.");
1717
1718 PyDoc_STRVAR(chain_from_iterable_doc,
1719 "chain.from_iterable(iterable) --> chain object\n\
1720 \n\
1721 Alternate chain() contructor taking a single iterable argument\n\
1722 that evaluates lazily.");
1723
1724 static PyMethodDef chain_methods[] = {
1725 {"from_iterable", (PyCFunction) chain_new_from_iterable, METH_O | METH_CLASS,
1726 chain_from_iterable_doc},
1727 {NULL, NULL} /* sentinel */
1728 };
1729
1730 static PyTypeObject chain_type = {
1731 PyVarObject_HEAD_INIT(NULL, 0)
1732 "itertools.chain", /* tp_name */
1733 sizeof(chainobject), /* tp_basicsize */
1734 0, /* tp_itemsize */
1735 /* methods */
1736 (destructor)chain_dealloc, /* tp_dealloc */
1737 0, /* tp_print */
1738 0, /* tp_getattr */
1739 0, /* tp_setattr */
1740 0, /* tp_compare */
1741 0, /* tp_repr */
1742 0, /* tp_as_number */
1743 0, /* tp_as_sequence */
1744 0, /* tp_as_mapping */
1745 0, /* tp_hash */
1746 0, /* tp_call */
1747 0, /* tp_str */
1748 PyObject_GenericGetAttr, /* tp_getattro */
1749 0, /* tp_setattro */
1750 0, /* tp_as_buffer */
1751 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1752 Py_TPFLAGS_BASETYPE, /* tp_flags */
1753 chain_doc, /* tp_doc */
1754 (traverseproc)chain_traverse, /* tp_traverse */
1755 0, /* tp_clear */
1756 0, /* tp_richcompare */
1757 0, /* tp_weaklistoffset */
1758 PyObject_SelfIter, /* tp_iter */
1759 (iternextfunc)chain_next, /* tp_iternext */
1760 chain_methods, /* tp_methods */
1761 0, /* tp_members */
1762 0, /* tp_getset */
1763 0, /* tp_base */
1764 0, /* tp_dict */
1765 0, /* tp_descr_get */
1766 0, /* tp_descr_set */
1767 0, /* tp_dictoffset */
1768 0, /* tp_init */
1769 0, /* tp_alloc */
1770 chain_new, /* tp_new */
1771 PyObject_GC_Del, /* tp_free */
1772 };
1773
1774
1775 /* product object ************************************************************/
1776
1777 typedef struct {
1778 PyObject_HEAD
1779 PyObject *pools; /* tuple of pool tuples */
1780 Py_ssize_t *indices; /* one index per pool */
1781 PyObject *result; /* most recently returned result tuple */
1782 int stopped; /* set to 1 when the product iterator is exhausted */
1783 } productobject;
1784
1785 static PyTypeObject product_type;
1786
1787 static PyObject *
1788 product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1789 {
1790 productobject *lz;
1791 Py_ssize_t nargs, npools, repeat=1;
1792 PyObject *pools = NULL;
1793 Py_ssize_t *indices = NULL;
1794 Py_ssize_t i;
1795
1796 if (kwds != NULL) {
1797 char *kwlist[] = {"repeat", 0};
1798 PyObject *tmpargs = PyTuple_New(0);
1799 if (tmpargs == NULL)
1800 return NULL;
1801 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product", kwlist, &repeat)) {
1802 Py_DECREF(tmpargs);
1803 return NULL;
1804 }
1805 Py_DECREF(tmpargs);
1806 if (repeat < 0) {
1807 PyErr_SetString(PyExc_ValueError,
1808 "repeat argument cannot be negative");
1809 return NULL;
1810 }
1811 }
1812
1813 assert(PyTuple_Check(args));
1814 nargs = (repeat == 0) ? 0 : PyTuple_GET_SIZE(args);
1815 npools = nargs * repeat;
1816
1817 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1818 if (indices == NULL) {
1819 PyErr_NoMemory();
1820 goto error;
1821 }
1822
1823 pools = PyTuple_New(npools);
1824 if (pools == NULL)
1825 goto error;
1826
1827 for (i=0; i < nargs ; ++i) {
1828 PyObject *item = PyTuple_GET_ITEM(args, i);
1829 PyObject *pool = PySequence_Tuple(item);
1830 if (pool == NULL)
1831 goto error;
1832 PyTuple_SET_ITEM(pools, i, pool);
1833 indices[i] = 0;
1834 }
1835 for ( ; i < npools; ++i) {
1836 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
1837 Py_INCREF(pool);
1838 PyTuple_SET_ITEM(pools, i, pool);
1839 indices[i] = 0;
1840 }
1841
1842 /* create productobject structure */
1843 lz = (productobject *)type->tp_alloc(type, 0);
1844 if (lz == NULL)
1845 goto error;
1846
1847 lz->pools = pools;
1848 lz->indices = indices;
1849 lz->result = NULL;
1850 lz->stopped = 0;
1851
1852 return (PyObject *)lz;
1853
1854 error:
1855 if (indices != NULL)
1856 PyMem_Free(indices);
1857 Py_XDECREF(pools);
1858 return NULL;
1859 }
1860
1861 static void
1862 product_dealloc(productobject *lz)
1863 {
1864 PyObject_GC_UnTrack(lz);
1865 Py_XDECREF(lz->pools);
1866 Py_XDECREF(lz->result);
1867 if (lz->indices != NULL)
1868 PyMem_Free(lz->indices);
1869 Py_TYPE(lz)->tp_free(lz);
1870 }
1871
1872 static int
1873 product_traverse(productobject *lz, visitproc visit, void *arg)
1874 {
1875 Py_VISIT(lz->pools);
1876 Py_VISIT(lz->result);
1877 return 0;
1878 }
1879
1880 static PyObject *
1881 product_next(productobject *lz)
1882 {
1883 PyObject *pool;
1884 PyObject *elem;
1885 PyObject *oldelem;
1886 PyObject *pools = lz->pools;
1887 PyObject *result = lz->result;
1888 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1889 Py_ssize_t i;
1890
1891 if (lz->stopped)
1892 return NULL;
1893
1894 if (result == NULL) {
1895 /* On the first pass, return an initial tuple filled with the
1896 first element from each pool. */
1897 result = PyTuple_New(npools);
1898 if (result == NULL)
1899 goto empty;
1900 lz->result = result;
1901 for (i=0; i < npools; i++) {
1902 pool = PyTuple_GET_ITEM(pools, i);
1903 if (PyTuple_GET_SIZE(pool) == 0)
1904 goto empty;
1905 elem = PyTuple_GET_ITEM(pool, 0);
1906 Py_INCREF(elem);
1907 PyTuple_SET_ITEM(result, i, elem);
1908 }
1909 } else {
1910 Py_ssize_t *indices = lz->indices;
1911
1912 /* Copy the previous result tuple or re-use it if available */
1913 if (Py_REFCNT(result) > 1) {
1914 PyObject *old_result = result;
1915 result = PyTuple_New(npools);
1916 if (result == NULL)
1917 goto empty;
1918 lz->result = result;
1919 for (i=0; i < npools; i++) {
1920 elem = PyTuple_GET_ITEM(old_result, i);
1921 Py_INCREF(elem);
1922 PyTuple_SET_ITEM(result, i, elem);
1923 }
1924 Py_DECREF(old_result);
1925 }
1926 /* Now, we've got the only copy so we can update it in-place */
1927 assert (npools==0 || Py_REFCNT(result) == 1);
1928
1929 /* Update the pool indices right-to-left. Only advance to the
1930 next pool when the previous one rolls-over */
1931 for (i=npools-1 ; i >= 0 ; i--) {
1932 pool = PyTuple_GET_ITEM(pools, i);
1933 indices[i]++;
1934 if (indices[i] == PyTuple_GET_SIZE(pool)) {
1935 /* Roll-over and advance to next pool */
1936 indices[i] = 0;
1937 elem = PyTuple_GET_ITEM(pool, 0);
1938 Py_INCREF(elem);
1939 oldelem = PyTuple_GET_ITEM(result, i);
1940 PyTuple_SET_ITEM(result, i, elem);
1941 Py_DECREF(oldelem);
1942 } else {
1943 /* No rollover. Just increment and stop here. */
1944 elem = PyTuple_GET_ITEM(pool, indices[i]);
1945 Py_INCREF(elem);
1946 oldelem = PyTuple_GET_ITEM(result, i);
1947 PyTuple_SET_ITEM(result, i, elem);
1948 Py_DECREF(oldelem);
1949 break;
1950 }
1951 }
1952
1953 /* If i is negative, then the indices have all rolled-over
1954 and we're done. */
1955 if (i < 0)
1956 goto empty;
1957 }
1958
1959 Py_INCREF(result);
1960 return result;
1961
1962 empty:
1963 lz->stopped = 1;
1964 return NULL;
1965 }
1966
1967 PyDoc_STRVAR(product_doc,
1968 "product(*iterables) --> product object\n\
1969 \n\
1970 Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
1971 For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1972 The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1973 cycle in a manner similar to an odometer (with the rightmost element changing\n\
1974 on every iteration).\n\n\
1975 To compute the product of an iterable with itself, specify the number\n\
1976 of repetitions with the optional repeat keyword argument. For example,\n\
1977 product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
1978 product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1979 product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1980
1981 static PyTypeObject product_type = {
1982 PyVarObject_HEAD_INIT(NULL, 0)
1983 "itertools.product", /* tp_name */
1984 sizeof(productobject), /* tp_basicsize */
1985 0, /* tp_itemsize */
1986 /* methods */
1987 (destructor)product_dealloc, /* tp_dealloc */
1988 0, /* tp_print */
1989 0, /* tp_getattr */
1990 0, /* tp_setattr */
1991 0, /* tp_compare */
1992 0, /* tp_repr */
1993 0, /* tp_as_number */
1994 0, /* tp_as_sequence */
1995 0, /* tp_as_mapping */
1996 0, /* tp_hash */
1997 0, /* tp_call */
1998 0, /* tp_str */
1999 PyObject_GenericGetAttr, /* tp_getattro */
2000 0, /* tp_setattro */
2001 0, /* tp_as_buffer */
2002 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2003 Py_TPFLAGS_BASETYPE, /* tp_flags */
2004 product_doc, /* tp_doc */
2005 (traverseproc)product_traverse, /* tp_traverse */
2006 0, /* tp_clear */
2007 0, /* tp_richcompare */
2008 0, /* tp_weaklistoffset */
2009 PyObject_SelfIter, /* tp_iter */
2010 (iternextfunc)product_next, /* tp_iternext */
2011 0, /* tp_methods */
2012 0, /* tp_members */
2013 0, /* tp_getset */
2014 0, /* tp_base */
2015 0, /* tp_dict */
2016 0, /* tp_descr_get */
2017 0, /* tp_descr_set */
2018 0, /* tp_dictoffset */
2019 0, /* tp_init */
2020 0, /* tp_alloc */
2021 product_new, /* tp_new */
2022 PyObject_GC_Del, /* tp_free */
2023 };
2024
2025
2026 /* combinations object ************************************************************/
2027
2028 typedef struct {
2029 PyObject_HEAD
2030 PyObject *pool; /* input converted to a tuple */
2031 Py_ssize_t *indices; /* one index per result element */
2032 PyObject *result; /* most recently returned result tuple */
2033 Py_ssize_t r; /* size of result tuple */
2034 int stopped; /* set to 1 when the combinations iterator is exhausted */
2035 } combinationsobject;
2036
2037 static PyTypeObject combinations_type;
2038
2039 static PyObject *
2040 combinations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2041 {
2042 combinationsobject *co;
2043 Py_ssize_t n;
2044 Py_ssize_t r;
2045 PyObject *pool = NULL;
2046 PyObject *iterable = NULL;
2047 Py_ssize_t *indices = NULL;
2048 Py_ssize_t i;
2049 static char *kwargs[] = {"iterable", "r", NULL};
2050
2051 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations", kwargs,
2052 &iterable, &r))
2053 return NULL;
2054
2055 pool = PySequence_Tuple(iterable);
2056 if (pool == NULL)
2057 goto error;
2058 n = PyTuple_GET_SIZE(pool);
2059 if (r < 0) {
2060 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2061 goto error;
2062 }
2063
2064 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2065 if (indices == NULL) {
2066 PyErr_NoMemory();
2067 goto error;
2068 }
2069
2070 for (i=0 ; i<r ; i++)
2071 indices[i] = i;
2072
2073 /* create combinationsobject structure */
2074 co = (combinationsobject *)type->tp_alloc(type, 0);
2075 if (co == NULL)
2076 goto error;
2077
2078 co->pool = pool;
2079 co->indices = indices;
2080 co->result = NULL;
2081 co->r = r;
2082 co->stopped = r > n ? 1 : 0;
2083
2084 return (PyObject *)co;
2085
2086 error:
2087 if (indices != NULL)
2088 PyMem_Free(indices);
2089 Py_XDECREF(pool);
2090 return NULL;
2091 }
2092
2093 static void
2094 combinations_dealloc(combinationsobject *co)
2095 {
2096 PyObject_GC_UnTrack(co);
2097 Py_XDECREF(co->pool);
2098 Py_XDECREF(co->result);
2099 if (co->indices != NULL)
2100 PyMem_Free(co->indices);
2101 Py_TYPE(co)->tp_free(co);
2102 }
2103
2104 static int
2105 combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2106 {
2107 Py_VISIT(co->pool);
2108 Py_VISIT(co->result);
2109 return 0;
2110 }
2111
2112 static PyObject *
2113 combinations_next(combinationsobject *co)
2114 {
2115 PyObject *elem;
2116 PyObject *oldelem;
2117 PyObject *pool = co->pool;
2118 Py_ssize_t *indices = co->indices;
2119 PyObject *result = co->result;
2120 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2121 Py_ssize_t r = co->r;
2122 Py_ssize_t i, j, index;
2123
2124 if (co->stopped)
2125 return NULL;
2126
2127 if (result == NULL) {
2128 /* On the first pass, initialize result tuple using the indices */
2129 result = PyTuple_New(r);
2130 if (result == NULL)
2131 goto empty;
2132 co->result = result;
2133 for (i=0; i<r ; i++) {
2134 index = indices[i];
2135 elem = PyTuple_GET_ITEM(pool, index);
2136 Py_INCREF(elem);
2137 PyTuple_SET_ITEM(result, i, elem);
2138 }
2139 } else {
2140 /* Copy the previous result tuple or re-use it if available */
2141 if (Py_REFCNT(result) > 1) {
2142 PyObject *old_result = result;
2143 result = PyTuple_New(r);
2144 if (result == NULL)
2145 goto empty;
2146 co->result = result;
2147 for (i=0; i<r ; i++) {
2148 elem = PyTuple_GET_ITEM(old_result, i);
2149 Py_INCREF(elem);
2150 PyTuple_SET_ITEM(result, i, elem);
2151 }
2152 Py_DECREF(old_result);
2153 }
2154 /* Now, we've got the only copy so we can update it in-place
2155 * CPython's empty tuple is a singleton and cached in
2156 * PyTuple's freelist.
2157 */
2158 assert(r == 0 || Py_REFCNT(result) == 1);
2159
2160 /* Scan indices right-to-left until finding one that is not
2161 at its maximum (i + n - r). */
2162 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2163 ;
2164
2165 /* If i is negative, then the indices are all at
2166 their maximum value and we're done. */
2167 if (i < 0)
2168 goto empty;
2169
2170 /* Increment the current index which we know is not at its
2171 maximum. Then move back to the right setting each index
2172 to its lowest possible value (one higher than the index
2173 to its left -- this maintains the sort order invariant). */
2174 indices[i]++;
2175 for (j=i+1 ; j<r ; j++)
2176 indices[j] = indices[j-1] + 1;
2177
2178 /* Update the result tuple for the new indices
2179 starting with i, the leftmost index that changed */
2180 for ( ; i<r ; i++) {
2181 index = indices[i];
2182 elem = PyTuple_GET_ITEM(pool, index);
2183 Py_INCREF(elem);
2184 oldelem = PyTuple_GET_ITEM(result, i);
2185 PyTuple_SET_ITEM(result, i, elem);
2186 Py_DECREF(oldelem);
2187 }
2188 }
2189
2190 Py_INCREF(result);
2191 return result;
2192
2193 empty:
2194 co->stopped = 1;
2195 return NULL;
2196 }
2197
2198 PyDoc_STRVAR(combinations_doc,
2199 "combinations(iterable[, r]) --> combinations object\n\
2200 \n\
2201 Return successive r-length combinations of elements in the iterable.\n\n\
2202 combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)");
2203
2204 static PyTypeObject combinations_type = {
2205 PyVarObject_HEAD_INIT(NULL, 0)
2206 "itertools.combinations", /* tp_name */
2207 sizeof(combinationsobject), /* tp_basicsize */
2208 0, /* tp_itemsize */
2209 /* methods */
2210 (destructor)combinations_dealloc, /* tp_dealloc */
2211 0, /* tp_print */
2212 0, /* tp_getattr */
2213 0, /* tp_setattr */
2214 0, /* tp_compare */
2215 0, /* tp_repr */
2216 0, /* tp_as_number */
2217 0, /* tp_as_sequence */
2218 0, /* tp_as_mapping */
2219 0, /* tp_hash */
2220 0, /* tp_call */
2221 0, /* tp_str */
2222 PyObject_GenericGetAttr, /* tp_getattro */
2223 0, /* tp_setattro */
2224 0, /* tp_as_buffer */
2225 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2226 Py_TPFLAGS_BASETYPE, /* tp_flags */
2227 combinations_doc, /* tp_doc */
2228 (traverseproc)combinations_traverse, /* tp_traverse */
2229 0, /* tp_clear */
2230 0, /* tp_richcompare */
2231 0, /* tp_weaklistoffset */
2232 PyObject_SelfIter, /* tp_iter */
2233 (iternextfunc)combinations_next, /* tp_iternext */
2234 0, /* tp_methods */
2235 0, /* tp_members */
2236 0, /* tp_getset */
2237 0, /* tp_base */
2238 0, /* tp_dict */
2239 0, /* tp_descr_get */
2240 0, /* tp_descr_set */
2241 0, /* tp_dictoffset */
2242 0, /* tp_init */
2243 0, /* tp_alloc */
2244 combinations_new, /* tp_new */
2245 PyObject_GC_Del, /* tp_free */
2246 };
2247
2248
2249 /* combinations with replacement object *******************************************/
2250
2251 /* Equivalent to:
2252
2253 def combinations_with_replacement(iterable, r):
2254 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2255 # number items returned: (n+r-1)! / r! / (n-1)!
2256 pool = tuple(iterable)
2257 n = len(pool)
2258 indices = [0] * r
2259 yield tuple(pool[i] for i in indices)
2260 while 1:
2261 for i in reversed(range(r)):
2262 if indices[i] != n - 1:
2263 break
2264 else:
2265 return
2266 indices[i:] = [indices[i] + 1] * (r - i)
2267 yield tuple(pool[i] for i in indices)
2268
2269 def combinations_with_replacement2(iterable, r):
2270 'Alternate version that filters from product()'
2271 pool = tuple(iterable)
2272 n = len(pool)
2273 for indices in product(range(n), repeat=r):
2274 if sorted(indices) == list(indices):
2275 yield tuple(pool[i] for i in indices)
2276 */
2277 typedef struct {
2278 PyObject_HEAD
2279 PyObject *pool; /* input converted to a tuple */
2280 Py_ssize_t *indices; /* one index per result element */
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 cwr iterator is exhausted */
2284 } cwrobject;
2285
2286 static PyTypeObject cwr_type;
2287
2288 static PyObject *
2289 cwr_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2290 {
2291 cwrobject *co;
2292 Py_ssize_t n;
2293 Py_ssize_t r;
2294 PyObject *pool = NULL;
2295 PyObject *iterable = NULL;
2296 Py_ssize_t *indices = NULL;
2297 Py_ssize_t i;
2298 static char *kwargs[] = {"iterable", "r", NULL};
2299
2300 if (!PyArg_ParseTupleAndKeywords(args, kwds, "On:combinations_with_replacement", kwargs,
2301 &iterable, &r))
2302 return NULL;
2303
2304 pool = PySequence_Tuple(iterable);
2305 if (pool == NULL)
2306 goto error;
2307 n = PyTuple_GET_SIZE(pool);
2308 if (r < 0) {
2309 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2310 goto error;
2311 }
2312
2313 indices = PyMem_Malloc(r * sizeof(Py_ssize_t));
2314 if (indices == NULL) {
2315 PyErr_NoMemory();
2316 goto error;
2317 }
2318
2319 for (i=0 ; i<r ; i++)
2320 indices[i] = 0;
2321
2322 /* create cwrobject structure */
2323 co = (cwrobject *)type->tp_alloc(type, 0);
2324 if (co == NULL)
2325 goto error;
2326
2327 co->pool = pool;
2328 co->indices = indices;
2329 co->result = NULL;
2330 co->r = r;
2331 co->stopped = !n && r;
2332
2333 return (PyObject *)co;
2334
2335 error:
2336 if (indices != NULL)
2337 PyMem_Free(indices);
2338 Py_XDECREF(pool);
2339 return NULL;
2340 }
2341
2342 static void
2343 cwr_dealloc(cwrobject *co)
2344 {
2345 PyObject_GC_UnTrack(co);
2346 Py_XDECREF(co->pool);
2347 Py_XDECREF(co->result);
2348 if (co->indices != NULL)
2349 PyMem_Free(co->indices);
2350 Py_TYPE(co)->tp_free(co);
2351 }
2352
2353 static int
2354 cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2355 {
2356 Py_VISIT(co->pool);
2357 Py_VISIT(co->result);
2358 return 0;
2359 }
2360
2361 static PyObject *
2362 cwr_next(cwrobject *co)
2363 {
2364 PyObject *elem;
2365 PyObject *oldelem;
2366 PyObject *pool = co->pool;
2367 Py_ssize_t *indices = co->indices;
2368 PyObject *result = co->result;
2369 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2370 Py_ssize_t r = co->r;
2371 Py_ssize_t i, j, index;
2372
2373 if (co->stopped)
2374 return NULL;
2375
2376 if (result == NULL) {
2377 /* On the first pass, initialize result tuple using the indices */
2378 result = PyTuple_New(r);
2379 if (result == NULL)
2380 goto empty;
2381 co->result = result;
2382 for (i=0; i<r ; i++) {
2383 index = indices[i];
2384 elem = PyTuple_GET_ITEM(pool, index);
2385 Py_INCREF(elem);
2386 PyTuple_SET_ITEM(result, i, elem);
2387 }
2388 } else {
2389 /* Copy the previous result tuple or re-use it if available */
2390 if (Py_REFCNT(result) > 1) {
2391 PyObject *old_result = result;
2392 result = PyTuple_New(r);
2393 if (result == NULL)
2394 goto empty;
2395 co->result = result;
2396 for (i=0; i<r ; i++) {
2397 elem = PyTuple_GET_ITEM(old_result, i);
2398 Py_INCREF(elem);
2399 PyTuple_SET_ITEM(result, i, elem);
2400 }
2401 Py_DECREF(old_result);
2402 }
2403 /* Now, we've got the only copy so we can update it in-place CPython's
2404 empty tuple is a singleton and cached in PyTuple's freelist. */
2405 assert(r == 0 || Py_REFCNT(result) == 1);
2406
2407 /* Scan indices right-to-left until finding one that is not
2408 * at its maximum (n-1). */
2409 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2410 ;
2411
2412 /* If i is negative, then the indices are all at
2413 their maximum value and we're done. */
2414 if (i < 0)
2415 goto empty;
2416
2417 /* Increment the current index which we know is not at its
2418 maximum. Then set all to the right to the same value. */
2419 indices[i]++;
2420 for (j=i+1 ; j<r ; j++)
2421 indices[j] = indices[j-1];
2422
2423 /* Update the result tuple for the new indices
2424 starting with i, the leftmost index that changed */
2425 for ( ; i<r ; i++) {
2426 index = indices[i];
2427 elem = PyTuple_GET_ITEM(pool, index);
2428 Py_INCREF(elem);
2429 oldelem = PyTuple_GET_ITEM(result, i);
2430 PyTuple_SET_ITEM(result, i, elem);
2431 Py_DECREF(oldelem);
2432 }
2433 }
2434
2435 Py_INCREF(result);
2436 return result;
2437
2438 empty:
2439 co->stopped = 1;
2440 return NULL;
2441 }
2442
2443 PyDoc_STRVAR(cwr_doc,
2444 "combinations_with_replacement(iterable[, r]) --> combinations_with_replacement object\n\
2445 \n\
2446 Return successive r-length combinations of elements in the iterable\n\
2447 allowing individual elements to have successive repeats.\n\
2448 combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC");
2449
2450 static PyTypeObject cwr_type = {
2451 PyVarObject_HEAD_INIT(NULL, 0)
2452 "itertools.combinations_with_replacement", /* tp_name */
2453 sizeof(cwrobject), /* tp_basicsize */
2454 0, /* tp_itemsize */
2455 /* methods */
2456 (destructor)cwr_dealloc, /* tp_dealloc */
2457 0, /* tp_print */
2458 0, /* tp_getattr */
2459 0, /* tp_setattr */
2460 0, /* tp_compare */
2461 0, /* tp_repr */
2462 0, /* tp_as_number */
2463 0, /* tp_as_sequence */
2464 0, /* tp_as_mapping */
2465 0, /* tp_hash */
2466 0, /* tp_call */
2467 0, /* tp_str */
2468 PyObject_GenericGetAttr, /* tp_getattro */
2469 0, /* tp_setattro */
2470 0, /* tp_as_buffer */
2471 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2472 Py_TPFLAGS_BASETYPE, /* tp_flags */
2473 cwr_doc, /* tp_doc */
2474 (traverseproc)cwr_traverse, /* tp_traverse */
2475 0, /* tp_clear */
2476 0, /* tp_richcompare */
2477 0, /* tp_weaklistoffset */
2478 PyObject_SelfIter, /* tp_iter */
2479 (iternextfunc)cwr_next, /* tp_iternext */
2480 0, /* tp_methods */
2481 0, /* tp_members */
2482 0, /* tp_getset */
2483 0, /* tp_base */
2484 0, /* tp_dict */
2485 0, /* tp_descr_get */
2486 0, /* tp_descr_set */
2487 0, /* tp_dictoffset */
2488 0, /* tp_init */
2489 0, /* tp_alloc */
2490 cwr_new, /* tp_new */
2491 PyObject_GC_Del, /* tp_free */
2492 };
2493
2494
2495 /* permutations object ************************************************************
2496
2497 def permutations(iterable, r=None):
2498 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
2499 pool = tuple(iterable)
2500 n = len(pool)
2501 r = n if r is None else r
2502 indices = range(n)
2503 cycles = range(n-r+1, n+1)[::-1]
2504 yield tuple(pool[i] for i in indices[:r])
2505 while n:
2506 for i in reversed(range(r)):
2507 cycles[i] -= 1
2508 if cycles[i] == 0:
2509 indices[i:] = indices[i+1:] + indices[i:i+1]
2510 cycles[i] = n - i
2511 else:
2512 j = cycles[i]
2513 indices[i], indices[-j] = indices[-j], indices[i]
2514 yield tuple(pool[i] for i in indices[:r])
2515 break
2516 else:
2517 return
2518 */
2519
2520 typedef struct {
2521 PyObject_HEAD
2522 PyObject *pool; /* input converted to a tuple */
2523 Py_ssize_t *indices; /* one index per element in the pool */
2524 Py_ssize_t *cycles; /* one rollover counter per element in the result */
2525 PyObject *result; /* most recently returned result tuple */
2526 Py_ssize_t r; /* size of result tuple */
2527 int stopped; /* set to 1 when the permutations iterator is exhausted */
2528 } permutationsobject;
2529
2530 static PyTypeObject permutations_type;
2531
2532 static PyObject *
2533 permutations_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2534 {
2535 permutationsobject *po;
2536 Py_ssize_t n;
2537 Py_ssize_t r;
2538 PyObject *robj = Py_None;
2539 PyObject *pool = NULL;
2540 PyObject *iterable = NULL;
2541 Py_ssize_t *indices = NULL;
2542 Py_ssize_t *cycles = NULL;
2543 Py_ssize_t i;
2544 static char *kwargs[] = {"iterable", "r", NULL};
2545
2546 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:permutations", kwargs,
2547 &iterable, &robj))
2548 return NULL;
2549
2550 pool = PySequence_Tuple(iterable);
2551 if (pool == NULL)
2552 goto error;
2553 n = PyTuple_GET_SIZE(pool);
2554
2555 r = n;
2556 if (robj != Py_None) {
2557 r = PyInt_AsSsize_t(robj);
2558 if (r == -1 && PyErr_Occurred())
2559 goto error;
2560 }
2561 if (r < 0) {
2562 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2563 goto error;
2564 }
2565
2566 indices = PyMem_Malloc(n * sizeof(Py_ssize_t));
2567 cycles = PyMem_Malloc(r * sizeof(Py_ssize_t));
2568 if (indices == NULL || cycles == NULL) {
2569 PyErr_NoMemory();
2570 goto error;
2571 }
2572
2573 for (i=0 ; i<n ; i++)
2574 indices[i] = i;
2575 for (i=0 ; i<r ; i++)
2576 cycles[i] = n - i;
2577
2578 /* create permutationsobject structure */
2579 po = (permutationsobject *)type->tp_alloc(type, 0);
2580 if (po == NULL)
2581 goto error;
2582
2583 po->pool = pool;
2584 po->indices = indices;
2585 po->cycles = cycles;
2586 po->result = NULL;
2587 po->r = r;
2588 po->stopped = r > n ? 1 : 0;
2589
2590 return (PyObject *)po;
2591
2592 error:
2593 if (indices != NULL)
2594 PyMem_Free(indices);
2595 if (cycles != NULL)
2596 PyMem_Free(cycles);
2597 Py_XDECREF(pool);
2598 return NULL;
2599 }
2600
2601 static void
2602 permutations_dealloc(permutationsobject *po)
2603 {
2604 PyObject_GC_UnTrack(po);
2605 Py_XDECREF(po->pool);
2606 Py_XDECREF(po->result);
2607 PyMem_Free(po->indices);
2608 PyMem_Free(po->cycles);
2609 Py_TYPE(po)->tp_free(po);
2610 }
2611
2612 static int
2613 permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
2614 {
2615 Py_VISIT(po->pool);
2616 Py_VISIT(po->result);
2617 return 0;
2618 }
2619
2620 static PyObject *
2621 permutations_next(permutationsobject *po)
2622 {
2623 PyObject *elem;
2624 PyObject *oldelem;
2625 PyObject *pool = po->pool;
2626 Py_ssize_t *indices = po->indices;
2627 Py_ssize_t *cycles = po->cycles;
2628 PyObject *result = po->result;
2629 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2630 Py_ssize_t r = po->r;
2631 Py_ssize_t i, j, k, index;
2632
2633 if (po->stopped)
2634 return NULL;
2635
2636 if (result == NULL) {
2637 /* On the first pass, initialize result tuple using the indices */
2638 result = PyTuple_New(r);
2639 if (result == NULL)
2640 goto empty;
2641 po->result = result;
2642 for (i=0; i<r ; i++) {
2643 index = indices[i];
2644 elem = PyTuple_GET_ITEM(pool, index);
2645 Py_INCREF(elem);
2646 PyTuple_SET_ITEM(result, i, elem);
2647 }
2648 } else {
2649 if (n == 0)
2650 goto empty;
2651
2652 /* Copy the previous result tuple or re-use it if available */
2653 if (Py_REFCNT(result) > 1) {
2654 PyObject *old_result = result;
2655 result = PyTuple_New(r);
2656 if (result == NULL)
2657 goto empty;
2658 po->result = result;
2659 for (i=0; i<r ; i++) {
2660 elem = PyTuple_GET_ITEM(old_result, i);
2661 Py_INCREF(elem);
2662 PyTuple_SET_ITEM(result, i, elem);
2663 }
2664 Py_DECREF(old_result);
2665 }
2666 /* Now, we've got the only copy so we can update it in-place */
2667 assert(r == 0 || Py_REFCNT(result) == 1);
2668
2669 /* Decrement rightmost cycle, moving leftward upon zero rollover */
2670 for (i=r-1 ; i>=0 ; i--) {
2671 cycles[i] -= 1;
2672 if (cycles[i] == 0) {
2673 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
2674 index = indices[i];
2675 for (j=i ; j<n-1 ; j++)
2676 indices[j] = indices[j+1];
2677 indices[n-1] = index;
2678 cycles[i] = n - i;
2679 } else {
2680 j = cycles[i];
2681 index = indices[i];
2682 indices[i] = indices[n-j];
2683 indices[n-j] = index;
2684
2685 for (k=i; k<r ; k++) {
2686 /* start with i, the leftmost element that changed */
2687 /* yield tuple(pool[k] for k in indices[:r]) */
2688 index = indices[k];
2689 elem = PyTuple_GET_ITEM(pool, index);
2690 Py_INCREF(elem);
2691 oldelem = PyTuple_GET_ITEM(result, k);
2692 PyTuple_SET_ITEM(result, k, elem);
2693 Py_DECREF(oldelem);
2694 }
2695 break;
2696 }
2697 }
2698 /* If i is negative, then the cycles have all
2699 rolled-over and we're done. */
2700 if (i < 0)
2701 goto empty;
2702 }
2703 Py_INCREF(result);
2704 return result;
2705
2706 empty:
2707 po->stopped = 1;
2708 return NULL;
2709 }
2710
2711 PyDoc_STRVAR(permutations_doc,
2712 "permutations(iterable[, r]) --> permutations object\n\
2713 \n\
2714 Return successive r-length permutations of elements in the iterable.\n\n\
2715 permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)");
2716
2717 static PyTypeObject permutations_type = {
2718 PyVarObject_HEAD_INIT(NULL, 0)
2719 "itertools.permutations", /* tp_name */
2720 sizeof(permutationsobject), /* tp_basicsize */
2721 0, /* tp_itemsize */
2722 /* methods */
2723 (destructor)permutations_dealloc, /* tp_dealloc */
2724 0, /* tp_print */
2725 0, /* tp_getattr */
2726 0, /* tp_setattr */
2727 0, /* tp_compare */
2728 0, /* tp_repr */
2729 0, /* tp_as_number */
2730 0, /* tp_as_sequence */
2731 0, /* tp_as_mapping */
2732 0, /* tp_hash */
2733 0, /* tp_call */
2734 0, /* tp_str */
2735 PyObject_GenericGetAttr, /* tp_getattro */
2736 0, /* tp_setattro */
2737 0, /* tp_as_buffer */
2738 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2739 Py_TPFLAGS_BASETYPE, /* tp_flags */
2740 permutations_doc, /* tp_doc */
2741 (traverseproc)permutations_traverse, /* tp_traverse */
2742 0, /* tp_clear */
2743 0, /* tp_richcompare */
2744 0, /* tp_weaklistoffset */
2745 PyObject_SelfIter, /* tp_iter */
2746 (iternextfunc)permutations_next, /* tp_iternext */
2747 0, /* tp_methods */
2748 0, /* tp_members */
2749 0, /* tp_getset */
2750 0, /* tp_base */
2751 0, /* tp_dict */
2752 0, /* tp_descr_get */
2753 0, /* tp_descr_set */
2754 0, /* tp_dictoffset */
2755 0, /* tp_init */
2756 0, /* tp_alloc */
2757 permutations_new, /* tp_new */
2758 PyObject_GC_Del, /* tp_free */
2759 };
2760
2761
2762 /* compress object ************************************************************/
2763
2764 /* Equivalent to:
2765
2766 def compress(data, selectors):
2767 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2768 return (d for d, s in izip(data, selectors) if s)
2769 */
2770
2771 typedef struct {
2772 PyObject_HEAD
2773 PyObject *data;
2774 PyObject *selectors;
2775 } compressobject;
2776
2777 static PyTypeObject compress_type;
2778
2779 static PyObject *
2780 compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2781 {
2782 PyObject *seq1, *seq2;
2783 PyObject *data=NULL, *selectors=NULL;
2784 compressobject *lz;
2785 static char *kwargs[] = {"data", "selectors", NULL};
2786
2787 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO:compress", kwargs, &seq1, &seq2))
2788 return NULL;
2789
2790 data = PyObject_GetIter(seq1);
2791 if (data == NULL)
2792 goto fail;
2793 selectors = PyObject_GetIter(seq2);
2794 if (selectors == NULL)
2795 goto fail;
2796
2797 /* create compressobject structure */
2798 lz = (compressobject *)type->tp_alloc(type, 0);
2799 if (lz == NULL)
2800 goto fail;
2801 lz->data = data;
2802 lz->selectors = selectors;
2803 return (PyObject *)lz;
2804
2805 fail:
2806 Py_XDECREF(data);
2807 Py_XDECREF(selectors);
2808 return NULL;
2809 }
2810
2811 static void
2812 compress_dealloc(compressobject *lz)
2813 {
2814 PyObject_GC_UnTrack(lz);
2815 Py_XDECREF(lz->data);
2816 Py_XDECREF(lz->selectors);
2817 Py_TYPE(lz)->tp_free(lz);
2818 }
2819
2820 static int
2821 compress_traverse(compressobject *lz, visitproc visit, void *arg)
2822 {
2823 Py_VISIT(lz->data);
2824 Py_VISIT(lz->selectors);
2825 return 0;
2826 }
2827
2828 static PyObject *
2829 compress_next(compressobject *lz)
2830 {
2831 PyObject *data = lz->data, *selectors = lz->selectors;
2832 PyObject *datum, *selector;
2833 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2834 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2835 int ok;
2836
2837 while (1) {
2838 /* Steps: get datum, get selector, evaluate selector.
2839 Order is important (to match the pure python version
2840 in terms of which input gets a chance to raise an
2841 exception first).
2842 */
2843
2844 datum = datanext(data);
2845 if (datum == NULL)
2846 return NULL;
2847
2848 selector = selectornext(selectors);
2849 if (selector == NULL) {
2850 Py_DECREF(datum);
2851 return NULL;
2852 }
2853
2854 ok = PyObject_IsTrue(selector);
2855 Py_DECREF(selector);
2856 if (ok == 1)
2857 return datum;
2858 Py_DECREF(datum);
2859 if (ok == -1)
2860 return NULL;
2861 }
2862 }
2863
2864 PyDoc_STRVAR(compress_doc,
2865 "compress(data, selectors) --> iterator over selected data\n\
2866 \n\
2867 Return data elements corresponding to true selector elements.\n\
2868 Forms a shorter iterator from selected data elements using the\n\
2869 selectors to choose the data elements.");
2870
2871 static PyTypeObject compress_type = {
2872 PyVarObject_HEAD_INIT(NULL, 0)
2873 "itertools.compress", /* tp_name */
2874 sizeof(compressobject), /* tp_basicsize */
2875 0, /* tp_itemsize */
2876 /* methods */
2877 (destructor)compress_dealloc, /* tp_dealloc */
2878 0, /* tp_print */
2879 0, /* tp_getattr */
2880 0, /* tp_setattr */
2881 0, /* tp_compare */
2882 0, /* tp_repr */
2883 0, /* tp_as_number */
2884 0, /* tp_as_sequence */
2885 0, /* tp_as_mapping */
2886 0, /* tp_hash */
2887 0, /* tp_call */
2888 0, /* tp_str */
2889 PyObject_GenericGetAttr, /* tp_getattro */
2890 0, /* tp_setattro */
2891 0, /* tp_as_buffer */
2892 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2893 Py_TPFLAGS_BASETYPE, /* tp_flags */
2894 compress_doc, /* tp_doc */
2895 (traverseproc)compress_traverse, /* tp_traverse */
2896 0, /* tp_clear */
2897 0, /* tp_richcompare */
2898 0, /* tp_weaklistoffset */
2899 PyObject_SelfIter, /* tp_iter */
2900 (iternextfunc)compress_next, /* tp_iternext */
2901 0, /* tp_methods */
2902 0, /* tp_members */
2903 0, /* tp_getset */
2904 0, /* tp_base */
2905 0, /* tp_dict */
2906 0, /* tp_descr_get */
2907 0, /* tp_descr_set */
2908 0, /* tp_dictoffset */
2909 0, /* tp_init */
2910 0, /* tp_alloc */
2911 compress_new, /* tp_new */
2912 PyObject_GC_Del, /* tp_free */
2913 };
2914
2915
2916 /* ifilter object ************************************************************/
2917
2918 typedef struct {
2919 PyObject_HEAD
2920 PyObject *func;
2921 PyObject *it;
2922 } ifilterobject;
2923
2924 static PyTypeObject ifilter_type;
2925
2926 static PyObject *
2927 ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2928 {
2929 PyObject *func, *seq;
2930 PyObject *it;
2931 ifilterobject *lz;
2932
2933 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
2934 return NULL;
2935
2936 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
2937 return NULL;
2938
2939 /* Get iterator. */
2940 it = PyObject_GetIter(seq);
2941 if (it == NULL)
2942 return NULL;
2943
2944 /* create ifilterobject structure */
2945 lz = (ifilterobject *)type->tp_alloc(type, 0);
2946 if (lz == NULL) {
2947 Py_DECREF(it);
2948 return NULL;
2949 }
2950 Py_INCREF(func);
2951 lz->func = func;
2952 lz->it = it;
2953
2954 return (PyObject *)lz;
2955 }
2956
2957 static void
2958 ifilter_dealloc(ifilterobject *lz)
2959 {
2960 PyObject_GC_UnTrack(lz);
2961 Py_XDECREF(lz->func);
2962 Py_XDECREF(lz->it);
2963 Py_TYPE(lz)->tp_free(lz);
2964 }
2965
2966 static int
2967 ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2968 {
2969 Py_VISIT(lz->it);
2970 Py_VISIT(lz->func);
2971 return 0;
2972 }
2973
2974 static PyObject *
2975 ifilter_next(ifilterobject *lz)
2976 {
2977 PyObject *item;
2978 PyObject *it = lz->it;
2979 long ok;
2980 PyObject *(*iternext)(PyObject *);
2981
2982 iternext = *Py_TYPE(it)->tp_iternext;
2983 for (;;) {
2984 item = iternext(it);
2985 if (item == NULL)
2986 return NULL;
2987
2988 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
2989 ok = PyObject_IsTrue(item);
2990 } else {
2991 PyObject *good;
2992 good = PyObject_CallFunctionObjArgs(lz->func,
2993 item, NULL);
2994 if (good == NULL) {
2995 Py_DECREF(item);
2996 return NULL;
2997 }
2998 ok = PyObject_IsTrue(good);
2999 Py_DECREF(good);
3000 }
3001 if (ok)
3002 return item;
3003 Py_DECREF(item);
3004 }
3005 }
3006
3007 PyDoc_STRVAR(ifilter_doc,
3008 "ifilter(function or None, sequence) --> ifilter object\n\
3009 \n\
3010 Return those items of sequence for which function(item) is true.\n\
3011 If function is None, return the items that are true.");
3012
3013 static PyTypeObject ifilter_type = {
3014 PyVarObject_HEAD_INIT(NULL, 0)
3015 "itertools.ifilter", /* tp_name */
3016 sizeof(ifilterobject), /* tp_basicsize */
3017 0, /* tp_itemsize */
3018 /* methods */
3019 (destructor)ifilter_dealloc, /* tp_dealloc */
3020 0, /* tp_print */
3021 0, /* tp_getattr */
3022 0, /* tp_setattr */
3023 0, /* tp_compare */
3024 0, /* tp_repr */
3025 0, /* tp_as_number */
3026 0, /* tp_as_sequence */
3027 0, /* tp_as_mapping */
3028 0, /* tp_hash */
3029 0, /* tp_call */
3030 0, /* tp_str */
3031 PyObject_GenericGetAttr, /* tp_getattro */
3032 0, /* tp_setattro */
3033 0, /* tp_as_buffer */
3034 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3035 Py_TPFLAGS_BASETYPE, /* tp_flags */
3036 ifilter_doc, /* tp_doc */
3037 (traverseproc)ifilter_traverse, /* tp_traverse */
3038 0, /* tp_clear */
3039 0, /* tp_richcompare */
3040 0, /* tp_weaklistoffset */
3041 PyObject_SelfIter, /* tp_iter */
3042 (iternextfunc)ifilter_next, /* tp_iternext */
3043 0, /* tp_methods */
3044 0, /* tp_members */
3045 0, /* tp_getset */
3046 0, /* tp_base */
3047 0, /* tp_dict */
3048 0, /* tp_descr_get */
3049 0, /* tp_descr_set */
3050 0, /* tp_dictoffset */
3051 0, /* tp_init */
3052 0, /* tp_alloc */
3053 ifilter_new, /* tp_new */
3054 PyObject_GC_Del, /* tp_free */
3055 };
3056
3057
3058 /* ifilterfalse object ************************************************************/
3059
3060 typedef struct {
3061 PyObject_HEAD
3062 PyObject *func;
3063 PyObject *it;
3064 } ifilterfalseobject;
3065
3066 static PyTypeObject ifilterfalse_type;
3067
3068 static PyObject *
3069 ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3070 {
3071 PyObject *func, *seq;
3072 PyObject *it;
3073 ifilterfalseobject *lz;
3074
3075 if (type == &ifilterfalse_type &&
3076 !_PyArg_NoKeywords("ifilterfalse()", kwds))
3077 return NULL;
3078
3079 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
3080 return NULL;
3081
3082 /* Get iterator. */
3083 it = PyObject_GetIter(seq);
3084 if (it == NULL)
3085 return NULL;
3086
3087 /* create ifilterfalseobject structure */
3088 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
3089 if (lz == NULL) {
3090 Py_DECREF(it);
3091 return NULL;
3092 }
3093 Py_INCREF(func);
3094 lz->func = func;
3095 lz->it = it;
3096
3097 return (PyObject *)lz;
3098 }
3099
3100 static void
3101 ifilterfalse_dealloc(ifilterfalseobject *lz)
3102 {
3103 PyObject_GC_UnTrack(lz);
3104 Py_XDECREF(lz->func);
3105 Py_XDECREF(lz->it);
3106 Py_TYPE(lz)->tp_free(lz);
3107 }
3108
3109 static int
3110 ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
3111 {
3112 Py_VISIT(lz->it);
3113 Py_VISIT(lz->func);
3114 return 0;
3115 }
3116
3117 static PyObject *
3118 ifilterfalse_next(ifilterfalseobject *lz)
3119 {
3120 PyObject *item;
3121 PyObject *it = lz->it;
3122 long ok;
3123 PyObject *(*iternext)(PyObject *);
3124
3125 iternext = *Py_TYPE(it)->tp_iternext;
3126 for (;;) {
3127 item = iternext(it);
3128 if (item == NULL)
3129 return NULL;
3130
3131 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3132 ok = PyObject_IsTrue(item);
3133 } else {
3134 PyObject *good;
3135 good = PyObject_CallFunctionObjArgs(lz->func,
3136 item, NULL);
3137 if (good == NULL) {
3138 Py_DECREF(item);
3139 return NULL;
3140 }
3141 ok = PyObject_IsTrue(good);
3142 Py_DECREF(good);
3143 }
3144 if (!ok)
3145 return item;
3146 Py_DECREF(item);
3147 }
3148 }
3149
3150 PyDoc_STRVAR(ifilterfalse_doc,
3151 "ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
3152 \n\
3153 Return those items of sequence for which function(item) is false.\n\
3154 If function is None, return the items that are false.");
3155
3156 static PyTypeObject ifilterfalse_type = {
3157 PyVarObject_HEAD_INIT(NULL, 0)
3158 "itertools.ifilterfalse", /* tp_name */
3159 sizeof(ifilterfalseobject), /* tp_basicsize */
3160 0, /* tp_itemsize */
3161 /* methods */
3162 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
3163 0, /* tp_print */
3164 0, /* tp_getattr */
3165 0, /* tp_setattr */
3166 0, /* tp_compare */
3167 0, /* tp_repr */
3168 0, /* tp_as_number */
3169 0, /* tp_as_sequence */
3170 0, /* tp_as_mapping */
3171 0, /* tp_hash */
3172 0, /* tp_call */
3173 0, /* tp_str */
3174 PyObject_GenericGetAttr, /* tp_getattro */
3175 0, /* tp_setattro */
3176 0, /* tp_as_buffer */
3177 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3178 Py_TPFLAGS_BASETYPE, /* tp_flags */
3179 ifilterfalse_doc, /* tp_doc */
3180 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
3181 0, /* tp_clear */
3182 0, /* tp_richcompare */
3183 0, /* tp_weaklistoffset */
3184 PyObject_SelfIter, /* tp_iter */
3185 (iternextfunc)ifilterfalse_next, /* tp_iternext */
3186 0, /* tp_methods */
3187 0, /* tp_members */
3188 0, /* tp_getset */
3189 0, /* tp_base */
3190 0, /* tp_dict */
3191 0, /* tp_descr_get */
3192 0, /* tp_descr_set */
3193 0, /* tp_dictoffset */
3194 0, /* tp_init */
3195 0, /* tp_alloc */
3196 ifilterfalse_new, /* tp_new */
3197 PyObject_GC_Del, /* tp_free */
3198 };
3199
3200
3201 /* count object ************************************************************/
3202
3203 typedef struct {
3204 PyObject_HEAD
3205 Py_ssize_t cnt;
3206 PyObject *long_cnt;
3207 PyObject *long_step;
3208 } countobject;
3209
3210 /* Counting logic and invariants:
3211
3212 fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3213
3214 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyInt(1));
3215 Advances with: cnt += 1
3216 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
3217
3218 slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3219
3220 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3221 All counting is done with python objects (no overflows or underflows).
3222 Advances with: long_cnt += long_step
3223 Step may be zero -- effectively a slow version of repeat(cnt).
3224 Either long_cnt or long_step may be a float, Fraction, or Decimal.
3225 */
3226
3227 static PyTypeObject count_type;
3228
3229 static PyObject *
3230 count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3231 {
3232 countobject *lz;
3233 int slow_mode = 0;
3234 Py_ssize_t cnt = 0;
3235 PyObject *long_cnt = NULL;
3236 PyObject *long_step = NULL;
3237 static char *kwlist[] = {"start", "step", 0};
3238
3239 if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO:count",
3240 kwlist, &long_cnt, &long_step))
3241 return NULL;
3242
3243 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
3244 (long_step != NULL && !PyNumber_Check(long_step))) {
3245 PyErr_SetString(PyExc_TypeError, "a number is required");
3246 return NULL;
3247 }
3248
3249 if (long_cnt != NULL) {
3250 cnt = PyInt_AsSsize_t(long_cnt);
3251 if ((cnt == -1 && PyErr_Occurred()) || !PyInt_Check(long_cnt)) {
3252 PyErr_Clear();
3253 slow_mode = 1;
3254 }
3255 Py_INCREF(long_cnt);
3256 } else {
3257 cnt = 0;
3258 long_cnt = PyInt_FromLong(0);
3259 }
3260
3261 /* If not specified, step defaults to 1 */
3262 if (long_step == NULL) {
3263 long_step = PyInt_FromLong(1);
3264 if (long_step == NULL) {
3265 Py_DECREF(long_cnt);
3266 return NULL;
3267 }
3268 } else
3269 Py_INCREF(long_step);
3270
3271 assert(long_cnt != NULL && long_step != NULL);
3272
3273 /* Fast mode only works when the step is 1 */
3274 if (!PyInt_Check(long_step) ||
3275 PyInt_AS_LONG(long_step) != 1) {
3276 slow_mode = 1;
3277 }
3278
3279 if (slow_mode)
3280 cnt = PY_SSIZE_T_MAX;
3281 else
3282 Py_CLEAR(long_cnt);
3283
3284 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && !slow_mode) ||
3285 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && slow_mode));
3286 assert(slow_mode ||
3287 (PyInt_Check(long_step) && PyInt_AS_LONG(long_step) == 1));
3288
3289 /* create countobject structure */
3290 lz = (countobject *)type->tp_alloc(type, 0);
3291 if (lz == NULL) {
3292 Py_XDECREF(long_cnt);
3293 return NULL;
3294 }
3295 lz->cnt = cnt;
3296 lz->long_cnt = long_cnt;
3297 lz->long_step = long_step;
3298
3299 return (PyObject *)lz;
3300 }
3301
3302 static void
3303 count_dealloc(countobject *lz)
3304 {
3305 PyObject_GC_UnTrack(lz);
3306 Py_XDECREF(lz->long_cnt);
3307 Py_XDECREF(lz->long_step);
3308 Py_TYPE(lz)->tp_free(lz);
3309 }
3310
3311 static int
3312 count_traverse(countobject *lz, visitproc visit, void *arg)
3313 {
3314 Py_VISIT(lz->long_cnt);
3315 Py_VISIT(lz->long_step);
3316 return 0;
3317 }
3318
3319 static PyObject *
3320 count_nextlong(countobject *lz)
3321 {
3322 PyObject *long_cnt;
3323 PyObject *stepped_up;
3324
3325 long_cnt = lz->long_cnt;
3326 if (long_cnt == NULL) {
3327 /* Switch to slow_mode */
3328 long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
3329 if (long_cnt == NULL)
3330 return NULL;
3331 }
3332 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
3333
3334 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
3335 if (stepped_up == NULL)
3336 return NULL;
3337 lz->long_cnt = stepped_up;
3338 return long_cnt;
3339 }
3340
3341 static PyObject *
3342 count_next(countobject *lz)
3343 {
3344 if (lz->cnt == PY_SSIZE_T_MAX)
3345 return count_nextlong(lz);
3346 return PyInt_FromSsize_t(lz->cnt++);
3347 }
3348
3349 static PyObject *
3350 count_repr(countobject *lz)
3351 {
3352 PyObject *cnt_repr, *step_repr = NULL;
3353 PyObject *result = NULL;
3354
3355 if (lz->cnt != PY_SSIZE_T_MAX)
3356 return PyString_FromFormat("count(%zd)", lz->cnt);
3357
3358 cnt_repr = PyObject_Repr(lz->long_cnt);
3359 if (cnt_repr == NULL)
3360 return NULL;
3361
3362 if (PyInt_Check(lz->long_step) && PyInt_AS_LONG(lz->long_step) == 1) {
3363 /* Don't display step when it is an integer equal to 1 */
3364 result = PyString_FromFormat("count(%s)",
3365 PyString_AS_STRING(cnt_repr));
3366 } else {
3367 step_repr = PyObject_Repr(lz->long_step);
3368 if (step_repr != NULL)
3369 result = PyString_FromFormat("count(%s, %s)",
3370 PyString_AS_STRING(cnt_repr),
3371 PyString_AS_STRING(step_repr));
3372 }
3373 Py_DECREF(cnt_repr);
3374 Py_XDECREF(step_repr);
3375 return result;
3376 }
3377
3378 PyDoc_STRVAR(count_doc,
3379 "count(start=0, step=1]) --> count object\n\
3380 \n\
3381 Return a count object whose .next() method returns consecutive values.\n\
3382 Equivalent to:\n\n\
3383 def count(firstval=0, step=1):\n\
3384 x = firstval\n\
3385 while 1:\n\
3386 yield x\n\
3387 x += step\n");
3388
3389 static PyTypeObject count_type = {
3390 PyVarObject_HEAD_INIT(NULL, 0)
3391 "itertools.count", /* tp_name */
3392 sizeof(countobject), /* tp_basicsize */
3393 0, /* tp_itemsize */
3394 /* methods */
3395 (destructor)count_dealloc, /* tp_dealloc */
3396 0, /* tp_print */
3397 0, /* tp_getattr */
3398 0, /* tp_setattr */
3399 0, /* tp_compare */
3400 (reprfunc)count_repr, /* tp_repr */
3401 0, /* tp_as_number */
3402 0, /* tp_as_sequence */
3403 0, /* tp_as_mapping */
3404 0, /* tp_hash */
3405 0, /* tp_call */
3406 0, /* tp_str */
3407 PyObject_GenericGetAttr, /* tp_getattro */
3408 0, /* tp_setattro */
3409 0, /* tp_as_buffer */
3410 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3411 Py_TPFLAGS_BASETYPE, /* tp_flags */
3412 count_doc, /* tp_doc */
3413 (traverseproc)count_traverse, /* tp_traverse */
3414 0, /* tp_clear */
3415 0, /* tp_richcompare */
3416 0, /* tp_weaklistoffset */
3417 PyObject_SelfIter, /* tp_iter */
3418 (iternextfunc)count_next, /* tp_iternext */
3419 0, /* tp_methods */
3420 0, /* tp_members */
3421 0, /* tp_getset */
3422 0, /* tp_base */
3423 0, /* tp_dict */
3424 0, /* tp_descr_get */
3425 0, /* tp_descr_set */
3426 0, /* tp_dictoffset */
3427 0, /* tp_init */
3428 0, /* tp_alloc */
3429 count_new, /* tp_new */
3430 PyObject_GC_Del, /* tp_free */
3431 };
3432
3433
3434 /* izip object ************************************************************/
3435
3436 #include "Python.h"
3437
3438 typedef struct {
3439 PyObject_HEAD
3440 Py_ssize_t tuplesize;
3441 PyObject *ittuple; /* tuple of iterators */
3442 PyObject *result;
3443 } izipobject;
3444
3445 static PyTypeObject izip_type;
3446
3447 static PyObject *
3448 izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3449 {
3450 izipobject *lz;
3451 Py_ssize_t i;
3452 PyObject *ittuple; /* tuple of iterators */
3453 PyObject *result;
3454 Py_ssize_t tuplesize = PySequence_Length(args);
3455
3456 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
3457 return NULL;
3458
3459 /* args must be a tuple */
3460 assert(PyTuple_Check(args));
3461
3462 /* obtain iterators */
3463 ittuple = PyTuple_New(tuplesize);
3464 if (ittuple == NULL)
3465 return NULL;
3466 for (i=0; i < tuplesize; ++i) {
3467 PyObject *item = PyTuple_GET_ITEM(args, i);
3468 PyObject *it = PyObject_GetIter(item);
3469 if (it == NULL) {
3470 if (PyErr_ExceptionMatches(PyExc_TypeError))
3471 PyErr_Format(PyExc_TypeError,
3472 "izip argument #%zd must support iteration",
3473 i+1);
3474 Py_DECREF(ittuple);
3475 return NULL;
3476 }
3477 PyTuple_SET_ITEM(ittuple, i, it);
3478 }
3479
3480 /* create a result holder */
3481 result = PyTuple_New(tuplesize);
3482 if (result == NULL) {
3483 Py_DECREF(ittuple);
3484 return NULL;
3485 }
3486 for (i=0 ; i < tuplesize ; i++) {
3487 Py_INCREF(Py_None);
3488 PyTuple_SET_ITEM(result, i, Py_None);
3489 }
3490
3491 /* create izipobject structure */
3492 lz = (izipobject *)type->tp_alloc(type, 0);
3493 if (lz == NULL) {
3494 Py_DECREF(ittuple);
3495 Py_DECREF(result);
3496 return NULL;
3497 }
3498 lz->ittuple = ittuple;
3499 lz->tuplesize = tuplesize;
3500 lz->result = result;
3501
3502 return (PyObject *)lz;
3503 }
3504
3505 static void
3506 izip_dealloc(izipobject *lz)
3507 {
3508 PyObject_GC_UnTrack(lz);
3509 Py_XDECREF(lz->ittuple);
3510 Py_XDECREF(lz->result);
3511 Py_TYPE(lz)->tp_free(lz);
3512 }
3513
3514 static int
3515 izip_traverse(izipobject *lz, visitproc visit, void *arg)
3516 {
3517 Py_VISIT(lz->ittuple);
3518 Py_VISIT(lz->result);
3519 return 0;
3520 }
3521
3522 static PyObject *
3523 izip_next(izipobject *lz)
3524 {
3525 Py_ssize_t i;
3526 Py_ssize_t tuplesize = lz->tuplesize;
3527 PyObject *result = lz->result;
3528 PyObject *it;
3529 PyObject *item;
3530 PyObject *olditem;
3531
3532 if (tuplesize == 0)
3533 return NULL;
3534 if (Py_REFCNT(result) == 1) {
3535 Py_INCREF(result);
3536 for (i=0 ; i < tuplesize ; i++) {
3537 it = PyTuple_GET_ITEM(lz->ittuple, i);
3538 item = (*Py_TYPE(it)->tp_iternext)(it);
3539 if (item == NULL) {
3540 Py_DECREF(result);
3541 return NULL;
3542 }
3543 olditem = PyTuple_GET_ITEM(result, i);
3544 PyTuple_SET_ITEM(result, i, item);
3545 Py_DECREF(olditem);
3546 }
3547 } else {
3548 result = PyTuple_New(tuplesize);
3549 if (result == NULL)
3550 return NULL;
3551 for (i=0 ; i < tuplesize ; i++) {
3552 it = PyTuple_GET_ITEM(lz->ittuple, i);
3553 item = (*Py_TYPE(it)->tp_iternext)(it);
3554 if (item == NULL) {
3555 Py_DECREF(result);
3556 return NULL;
3557 }
3558 PyTuple_SET_ITEM(result, i, item);
3559 }
3560 }
3561 return result;
3562 }
3563
3564 PyDoc_STRVAR(izip_doc,
3565 "izip(iter1 [,iter2 [...]]) --> izip object\n\
3566 \n\
3567 Return a izip object whose .next() method returns a tuple where\n\
3568 the i-th element comes from the i-th iterable argument. The .next()\n\
3569 method continues until the shortest iterable in the argument sequence\n\
3570 is exhausted and then it raises StopIteration. Works like the zip()\n\
3571 function but consumes less memory by returning an iterator instead of\n\
3572 a list.");
3573
3574 static PyTypeObject izip_type = {
3575 PyVarObject_HEAD_INIT(NULL, 0)
3576 "itertools.izip", /* tp_name */
3577 sizeof(izipobject), /* tp_basicsize */
3578 0, /* tp_itemsize */
3579 /* methods */
3580 (destructor)izip_dealloc, /* tp_dealloc */
3581 0, /* tp_print */
3582 0, /* tp_getattr */
3583 0, /* tp_setattr */
3584 0, /* tp_compare */
3585 0, /* tp_repr */
3586 0, /* tp_as_number */
3587 0, /* tp_as_sequence */
3588 0, /* tp_as_mapping */
3589 0, /* tp_hash */
3590 0, /* tp_call */
3591 0, /* tp_str */
3592 PyObject_GenericGetAttr, /* tp_getattro */
3593 0, /* tp_setattro */
3594 0, /* tp_as_buffer */
3595 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3596 Py_TPFLAGS_BASETYPE, /* tp_flags */
3597 izip_doc, /* tp_doc */
3598 (traverseproc)izip_traverse, /* tp_traverse */
3599 0, /* tp_clear */
3600 0, /* tp_richcompare */
3601 0, /* tp_weaklistoffset */
3602 PyObject_SelfIter, /* tp_iter */
3603 (iternextfunc)izip_next, /* tp_iternext */
3604 0, /* tp_methods */
3605 0, /* tp_members */
3606 0, /* tp_getset */
3607 0, /* tp_base */
3608 0, /* tp_dict */
3609 0, /* tp_descr_get */
3610 0, /* tp_descr_set */
3611 0, /* tp_dictoffset */
3612 0, /* tp_init */
3613 0, /* tp_alloc */
3614 izip_new, /* tp_new */
3615 PyObject_GC_Del, /* tp_free */
3616 };
3617
3618
3619 /* repeat object ************************************************************/
3620
3621 typedef struct {
3622 PyObject_HEAD
3623 PyObject *element;
3624 Py_ssize_t cnt;
3625 } repeatobject;
3626
3627 static PyTypeObject repeat_type;
3628
3629 static PyObject *
3630 repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3631 {
3632 repeatobject *ro;
3633 PyObject *element;
3634 Py_ssize_t cnt = -1;
3635 static char *kwargs[] = {"object", "times", NULL};
3636
3637 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
3638 &element, &cnt))
3639 return NULL;
3640
3641 if (PyTuple_Size(args) == 2 && cnt < 0)
3642 cnt = 0;
3643
3644 ro = (repeatobject *)type->tp_alloc(type, 0);
3645 if (ro == NULL)
3646 return NULL;
3647 Py_INCREF(element);
3648 ro->element = element;
3649 ro->cnt = cnt;
3650 return (PyObject *)ro;
3651 }
3652
3653 static void
3654 repeat_dealloc(repeatobject *ro)
3655 {
3656 PyObject_GC_UnTrack(ro);
3657 Py_XDECREF(ro->element);
3658 Py_TYPE(ro)->tp_free(ro);
3659 }
3660
3661 static int
3662 repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
3663 {
3664 Py_VISIT(ro->element);
3665 return 0;
3666 }
3667
3668 static PyObject *
3669 repeat_next(repeatobject *ro)
3670 {
3671 if (ro->cnt == 0)
3672 return NULL;
3673 if (ro->cnt > 0)
3674 ro->cnt--;
3675 Py_INCREF(ro->element);
3676 return ro->element;
3677 }
3678
3679 static PyObject *
3680 repeat_repr(repeatobject *ro)
3681 {
3682 PyObject *result, *objrepr;
3683
3684 objrepr = PyObject_Repr(ro->element);
3685 if (objrepr == NULL)
3686 return NULL;
3687
3688 if (ro->cnt == -1)
3689 result = PyString_FromFormat("repeat(%s)",
3690 PyString_AS_STRING(objrepr));
3691 else
3692 result = PyString_FromFormat("repeat(%s, %zd)",
3693 PyString_AS_STRING(objrepr), ro->cnt);
3694 Py_DECREF(objrepr);
3695 return result;
3696 }
3697
3698 static PyObject *
3699 repeat_len(repeatobject *ro)
3700 {
3701 if (ro->cnt == -1) {
3702 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
3703 return NULL;
3704 }
3705 return PyInt_FromSize_t(ro->cnt);
3706 }
3707
3708 PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
3709
3710 static PyMethodDef repeat_methods[] = {
3711 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
3712 {NULL, NULL} /* sentinel */
3713 };
3714
3715 PyDoc_STRVAR(repeat_doc,
3716 "repeat(object [,times]) -> create an iterator which returns the object\n\
3717 for the specified number of times. If not specified, returns the object\n\
3718 endlessly.");
3719
3720 static PyTypeObject repeat_type = {
3721 PyVarObject_HEAD_INIT(NULL, 0)
3722 "itertools.repeat", /* tp_name */
3723 sizeof(repeatobject), /* tp_basicsize */
3724 0, /* tp_itemsize */
3725 /* methods */
3726 (destructor)repeat_dealloc, /* tp_dealloc */
3727 0, /* tp_print */
3728 0, /* tp_getattr */
3729 0, /* tp_setattr */
3730 0, /* tp_compare */
3731 (reprfunc)repeat_repr, /* tp_repr */
3732 0, /* tp_as_number */
3733 0, /* tp_as_sequence */
3734 0, /* tp_as_mapping */
3735 0, /* tp_hash */
3736 0, /* tp_call */
3737 0, /* tp_str */
3738 PyObject_GenericGetAttr, /* tp_getattro */
3739 0, /* tp_setattro */
3740 0, /* tp_as_buffer */
3741 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3742 Py_TPFLAGS_BASETYPE, /* tp_flags */
3743 repeat_doc, /* tp_doc */
3744 (traverseproc)repeat_traverse, /* tp_traverse */
3745 0, /* tp_clear */
3746 0, /* tp_richcompare */
3747 0, /* tp_weaklistoffset */
3748 PyObject_SelfIter, /* tp_iter */
3749 (iternextfunc)repeat_next, /* tp_iternext */
3750 repeat_methods, /* tp_methods */
3751 0, /* tp_members */
3752 0, /* tp_getset */
3753 0, /* tp_base */
3754 0, /* tp_dict */
3755 0, /* tp_descr_get */
3756 0, /* tp_descr_set */
3757 0, /* tp_dictoffset */
3758 0, /* tp_init */
3759 0, /* tp_alloc */
3760 repeat_new, /* tp_new */
3761 PyObject_GC_Del, /* tp_free */
3762 };
3763
3764 /* iziplongest object ************************************************************/
3765
3766 #include "Python.h"
3767
3768 typedef struct {
3769 PyObject_HEAD
3770 Py_ssize_t tuplesize;
3771 Py_ssize_t numactive;
3772 PyObject *ittuple; /* tuple of iterators */
3773 PyObject *result;
3774 PyObject *fillvalue;
3775 } iziplongestobject;
3776
3777 static PyTypeObject iziplongest_type;
3778
3779 static PyObject *
3780 izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3781 {
3782 iziplongestobject *lz;
3783 Py_ssize_t i;
3784 PyObject *ittuple; /* tuple of iterators */
3785 PyObject *result;
3786 PyObject *fillvalue = Py_None;
3787 Py_ssize_t tuplesize = PySequence_Length(args);
3788
3789 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
3790 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
3791 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
3792 PyErr_SetString(PyExc_TypeError,
3793 "izip_longest() got an unexpected keyword argument");
3794 return NULL;
3795 }
3796 }
3797
3798 /* args must be a tuple */
3799 assert(PyTuple_Check(args));
3800
3801 /* obtain iterators */
3802 ittuple = PyTuple_New(tuplesize);
3803 if (ittuple == NULL)
3804 return NULL;
3805 for (i=0; i < tuplesize; ++i) {
3806 PyObject *item = PyTuple_GET_ITEM(args, i);
3807 PyObject *it = PyObject_GetIter(item);
3808 if (it == NULL) {
3809 if (PyErr_ExceptionMatches(PyExc_TypeError))
3810 PyErr_Format(PyExc_TypeError,
3811 "izip_longest argument #%zd must support iteration",
3812 i+1);
3813 Py_DECREF(ittuple);
3814 return NULL;
3815 }
3816 PyTuple_SET_ITEM(ittuple, i, it);
3817 }
3818
3819 /* create a result holder */
3820 result = PyTuple_New(tuplesize);
3821 if (result == NULL) {
3822 Py_DECREF(ittuple);
3823 return NULL;
3824 }
3825 for (i=0 ; i < tuplesize ; i++) {
3826 Py_INCREF(Py_None);
3827 PyTuple_SET_ITEM(result, i, Py_None);
3828 }
3829
3830 /* create iziplongestobject structure */
3831 lz = (iziplongestobject *)type->tp_alloc(type, 0);
3832 if (lz == NULL) {
3833 Py_DECREF(ittuple);
3834 Py_DECREF(result);
3835 return NULL;
3836 }
3837 lz->ittuple = ittuple;
3838 lz->tuplesize = tuplesize;
3839 lz->numactive = tuplesize;
3840 lz->result = result;
3841 Py_INCREF(fillvalue);
3842 lz->fillvalue = fillvalue;
3843 return (PyObject *)lz;
3844 }
3845
3846 static void
3847 izip_longest_dealloc(iziplongestobject *lz)
3848 {
3849 PyObject_GC_UnTrack(lz);
3850 Py_XDECREF(lz->ittuple);
3851 Py_XDECREF(lz->result);
3852 Py_XDECREF(lz->fillvalue);
3853 Py_TYPE(lz)->tp_free(lz);
3854 }
3855
3856 static int
3857 izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
3858 {
3859 Py_VISIT(lz->ittuple);
3860 Py_VISIT(lz->result);
3861 Py_VISIT(lz->fillvalue);
3862 return 0;
3863 }
3864
3865 static PyObject *
3866 izip_longest_next(iziplongestobject *lz)
3867 {
3868 Py_ssize_t i;
3869 Py_ssize_t tuplesize = lz->tuplesize;
3870 PyObject *result = lz->result;
3871 PyObject *it;
3872 PyObject *item;
3873 PyObject *olditem;
3874
3875 if (tuplesize == 0)
3876 return NULL;
3877 if (lz->numactive == 0)
3878 return NULL;
3879 if (Py_REFCNT(result) == 1) {
3880 Py_INCREF(result);
3881 for (i=0 ; i < tuplesize ; i++) {
3882 it = PyTuple_GET_ITEM(lz->ittuple, i);
3883 if (it == NULL) {
3884 Py_INCREF(lz->fillvalue);
3885 item = lz->fillvalue;
3886 } else {
3887 item = (*Py_TYPE(it)->tp_iternext)(it);
3888 if (item == NULL) {
3889 lz->numactive -= 1;
3890 if (lz->numactive == 0) {
3891 Py_DECREF(result);
3892 return NULL;
3893 } else {
3894 Py_INCREF(lz->fillvalue);
3895 item = lz->fillvalue;
3896 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3897 Py_DECREF(it);
3898 }
3899 }
3900 }
3901 olditem = PyTuple_GET_ITEM(result, i);
3902 PyTuple_SET_ITEM(result, i, item);
3903 Py_DECREF(olditem);
3904 }
3905 } else {
3906 result = PyTuple_New(tuplesize);
3907 if (result == NULL)
3908 return NULL;
3909 for (i=0 ; i < tuplesize ; i++) {
3910 it = PyTuple_GET_ITEM(lz->ittuple, i);
3911 if (it == NULL) {
3912 Py_INCREF(lz->fillvalue);
3913 item = lz->fillvalue;
3914 } else {
3915 item = (*Py_TYPE(it)->tp_iternext)(it);
3916 if (item == NULL) {
3917 lz->numactive -= 1;
3918 if (lz->numactive == 0) {
3919 Py_DECREF(result);
3920 return NULL;
3921 } else {
3922 Py_INCREF(lz->fillvalue);
3923 item = lz->fillvalue;
3924 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
3925 Py_DECREF(it);
3926 }
3927 }
3928 }
3929 PyTuple_SET_ITEM(result, i, item);
3930 }
3931 }
3932 return result;
3933 }
3934
3935 PyDoc_STRVAR(izip_longest_doc,
3936 "izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
3937 \n\
3938 Return an izip_longest object whose .next() method returns a tuple where\n\
3939 the i-th element comes from the i-th iterable argument. The .next()\n\
3940 method continues until the longest iterable in the argument sequence\n\
3941 is exhausted and then it raises StopIteration. When the shorter iterables\n\
3942 are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
3943 defaults to None or can be specified by a keyword argument.\n\
3944 ");
3945
3946 static PyTypeObject iziplongest_type = {
3947 PyVarObject_HEAD_INIT(NULL, 0)
3948 "itertools.izip_longest", /* tp_name */
3949 sizeof(iziplongestobject), /* tp_basicsize */
3950 0, /* tp_itemsize */
3951 /* methods */
3952 (destructor)izip_longest_dealloc, /* tp_dealloc */
3953 0, /* tp_print */
3954 0, /* tp_getattr */
3955 0, /* tp_setattr */
3956 0, /* tp_compare */
3957 0, /* tp_repr */
3958 0, /* tp_as_number */
3959 0, /* tp_as_sequence */
3960 0, /* tp_as_mapping */
3961 0, /* tp_hash */
3962 0, /* tp_call */
3963 0, /* tp_str */
3964 PyObject_GenericGetAttr, /* tp_getattro */
3965 0, /* tp_setattro */
3966 0, /* tp_as_buffer */
3967 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3968 Py_TPFLAGS_BASETYPE, /* tp_flags */
3969 izip_longest_doc, /* tp_doc */
3970 (traverseproc)izip_longest_traverse, /* tp_traverse */
3971 0, /* tp_clear */
3972 0, /* tp_richcompare */
3973 0, /* tp_weaklistoffset */
3974 PyObject_SelfIter, /* tp_iter */
3975 (iternextfunc)izip_longest_next, /* tp_iternext */
3976 0, /* tp_methods */
3977 0, /* tp_members */
3978 0, /* tp_getset */
3979 0, /* tp_base */
3980 0, /* tp_dict */
3981 0, /* tp_descr_get */
3982 0, /* tp_descr_set */
3983 0, /* tp_dictoffset */
3984 0, /* tp_init */
3985 0, /* tp_alloc */
3986 izip_longest_new, /* tp_new */
3987 PyObject_GC_Del, /* tp_free */
3988 };
3989
3990 /* module level code ********************************************************/
3991
3992 PyDoc_STRVAR(module_doc,
3993 "Functional tools for creating and using iterators.\n\
3994 \n\
3995 Infinite iterators:\n\
3996 count([n]) --> n, n+1, n+2, ...\n\
3997 cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
3998 repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
3999 \n\
4000 Iterators terminating on the shortest input sequence:\n\
4001 chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
4002 compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4003 dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4004 groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
4005 ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
4006 ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
4007 islice(seq, [start,] stop [, step]) --> elements from\n\
4008 seq[start:stop:step]\n\
4009 imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
4010 starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
4011 tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
4012 takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
4013 izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4014 izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
4015 \n\
4016 Combinatoric generators:\n\
4017 product(p, q, ... [repeat=1]) --> cartesian product\n\
4018 permutations(p[, r])\n\
4019 combinations(p[, r])\n\
4020 combinations_with_replacement(p[, r])\n\
4021 ");
4022
4023
4024 static PyMethodDef module_methods[] = {
4025 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
4026 {NULL, NULL} /* sentinel */
4027 };
4028
4029 PyMODINIT_FUNC
4030 inititertools(void)
4031 {
4032 int i;
4033 PyObject *m;
4034 char *name;
4035 PyTypeObject *typelist[] = {
4036 &combinations_type,
4037 &cwr_type,
4038 &cycle_type,
4039 &dropwhile_type,
4040 &takewhile_type,
4041 &islice_type,
4042 &starmap_type,
4043 &imap_type,
4044 &chain_type,
4045 &compress_type,
4046 &ifilter_type,
4047 &ifilterfalse_type,
4048 &count_type,
4049 &izip_type,
4050 &iziplongest_type,
4051 &permutations_type,
4052 &product_type,
4053 &repeat_type,
4054 &groupby_type,
4055 NULL
4056 };
4057
4058 Py_TYPE(&teedataobject_type) = &PyType_Type;
4059 m = Py_InitModule3("itertools", module_methods, module_doc);
4060 if (m == NULL)
4061 return;
4062
4063 for (i=0 ; typelist[i] != NULL ; i++) {
4064 if (PyType_Ready(typelist[i]) < 0)
4065 return;
4066 name = strchr(typelist[i]->tp_name, '.');
4067 assert (name != NULL);
4068 Py_INCREF(typelist[i]);
4069 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
4070 }
4071
4072 if (PyType_Ready(&teedataobject_type) < 0)
4073 return;
4074 if (PyType_Ready(&tee_type) < 0)
4075 return;
4076 if (PyType_Ready(&_grouper_type) < 0)
4077 return;
4078 }
python git mirror