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