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