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