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