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