blob: 21659fbca2e1606bcfb71fae0c6364f2814e5a84 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
3
4/* Itertools module written and maintained
5 by Raymond D. Hettinger <python@rcn.com>
6 Copyright (c) 2003 Python Software Foundation.
7 All rights reserved.
8*/
9
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000010
11/* groupby object ***********************************************************/
12
13typedef struct {
14 PyObject_HEAD
15 PyObject *it;
16 PyObject *keyfunc;
17 PyObject *tgtkey;
18 PyObject *currkey;
19 PyObject *currvalue;
20} groupbyobject;
21
22static PyTypeObject groupby_type;
23static PyObject *_grouper_create(groupbyobject *, PyObject *);
24
25static PyObject *
26groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
27{
28 static char *kwargs[] = {"iterable", "key", NULL};
29 groupbyobject *gbo;
30 PyObject *it, *keyfunc = Py_None;
31
32 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
33 &it, &keyfunc))
34 return NULL;
35
36 gbo = (groupbyobject *)type->tp_alloc(type, 0);
37 if (gbo == NULL)
38 return NULL;
39 gbo->tgtkey = NULL;
40 gbo->currkey = NULL;
41 gbo->currvalue = NULL;
42 gbo->keyfunc = keyfunc;
43 Py_INCREF(keyfunc);
44 gbo->it = PyObject_GetIter(it);
45 if (gbo->it == NULL) {
46 Py_DECREF(gbo);
47 return NULL;
48 }
49 return (PyObject *)gbo;
50}
51
52static void
53groupby_dealloc(groupbyobject *gbo)
54{
55 PyObject_GC_UnTrack(gbo);
56 Py_XDECREF(gbo->it);
57 Py_XDECREF(gbo->keyfunc);
58 Py_XDECREF(gbo->tgtkey);
59 Py_XDECREF(gbo->currkey);
60 Py_XDECREF(gbo->currvalue);
61 gbo->ob_type->tp_free(gbo);
62}
63
64static int
65groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
66{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +000067 Py_VISIT(gbo->it);
68 Py_VISIT(gbo->keyfunc);
69 Py_VISIT(gbo->tgtkey);
70 Py_VISIT(gbo->currkey);
71 Py_VISIT(gbo->currvalue);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000072 return 0;
73}
74
75static PyObject *
76groupby_next(groupbyobject *gbo)
77{
78 PyObject *newvalue, *newkey, *r, *grouper;
79
80 /* skip to next iteration group */
81 for (;;) {
82 if (gbo->currkey == NULL)
83 /* pass */;
84 else if (gbo->tgtkey == NULL)
85 break;
86 else {
87 int rcmp;
88
89 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
90 gbo->currkey, Py_EQ);
91 if (rcmp == -1)
92 return NULL;
93 else if (rcmp == 0)
94 break;
95 }
96
97 newvalue = PyIter_Next(gbo->it);
98 if (newvalue == NULL)
99 return NULL;
100
101 if (gbo->keyfunc == Py_None) {
102 newkey = newvalue;
103 Py_INCREF(newvalue);
104 } else {
105 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
106 newvalue, NULL);
107 if (newkey == NULL) {
108 Py_DECREF(newvalue);
109 return NULL;
110 }
111 }
112
113 Py_XDECREF(gbo->currkey);
114 gbo->currkey = newkey;
115 Py_XDECREF(gbo->currvalue);
116 gbo->currvalue = newvalue;
117 }
118
119 Py_XDECREF(gbo->tgtkey);
120 gbo->tgtkey = gbo->currkey;
121 Py_INCREF(gbo->currkey);
122
123 grouper = _grouper_create(gbo, gbo->tgtkey);
124 if (grouper == NULL)
125 return NULL;
126
127 r = PyTuple_Pack(2, gbo->currkey, grouper);
128 Py_DECREF(grouper);
129 return r;
130}
131
132PyDoc_STRVAR(groupby_doc,
133"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
134(key, sub-iterator) grouped by each value of key(value).\n");
135
136static PyTypeObject groupby_type = {
137 PyObject_HEAD_INIT(NULL)
138 0, /* ob_size */
139 "itertools.groupby", /* tp_name */
140 sizeof(groupbyobject), /* tp_basicsize */
141 0, /* tp_itemsize */
142 /* methods */
143 (destructor)groupby_dealloc, /* tp_dealloc */
144 0, /* tp_print */
145 0, /* tp_getattr */
146 0, /* tp_setattr */
147 0, /* tp_compare */
148 0, /* tp_repr */
149 0, /* tp_as_number */
150 0, /* tp_as_sequence */
151 0, /* tp_as_mapping */
152 0, /* tp_hash */
153 0, /* tp_call */
154 0, /* tp_str */
155 PyObject_GenericGetAttr, /* tp_getattro */
156 0, /* tp_setattro */
157 0, /* tp_as_buffer */
158 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
159 Py_TPFLAGS_BASETYPE, /* tp_flags */
160 groupby_doc, /* tp_doc */
161 (traverseproc)groupby_traverse, /* tp_traverse */
162 0, /* tp_clear */
163 0, /* tp_richcompare */
164 0, /* tp_weaklistoffset */
165 PyObject_SelfIter, /* tp_iter */
166 (iternextfunc)groupby_next, /* tp_iternext */
167 0, /* tp_methods */
168 0, /* tp_members */
169 0, /* tp_getset */
170 0, /* tp_base */
171 0, /* tp_dict */
172 0, /* tp_descr_get */
173 0, /* tp_descr_set */
174 0, /* tp_dictoffset */
175 0, /* tp_init */
176 0, /* tp_alloc */
177 groupby_new, /* tp_new */
178 PyObject_GC_Del, /* tp_free */
179};
180
181
182/* _grouper object (internal) ************************************************/
183
184typedef struct {
185 PyObject_HEAD
186 PyObject *parent;
187 PyObject *tgtkey;
188} _grouperobject;
189
190static PyTypeObject _grouper_type;
191
192static PyObject *
193_grouper_create(groupbyobject *parent, PyObject *tgtkey)
194{
195 _grouperobject *igo;
196
197 igo = PyObject_New(_grouperobject, &_grouper_type);
198 if (igo == NULL)
199 return NULL;
200 igo->parent = (PyObject *)parent;
201 Py_INCREF(parent);
202 igo->tgtkey = tgtkey;
203 Py_INCREF(tgtkey);
204
205 return (PyObject *)igo;
206}
207
208static void
209_grouper_dealloc(_grouperobject *igo)
210{
211 Py_DECREF(igo->parent);
212 Py_DECREF(igo->tgtkey);
213 PyObject_Del(igo);
214}
215
216static PyObject *
217_grouper_next(_grouperobject *igo)
218{
219 groupbyobject *gbo = (groupbyobject *)igo->parent;
220 PyObject *newvalue, *newkey, *r;
221 int rcmp;
222
223 if (gbo->currvalue == NULL) {
224 newvalue = PyIter_Next(gbo->it);
225 if (newvalue == NULL)
226 return NULL;
227
228 if (gbo->keyfunc == Py_None) {
229 newkey = newvalue;
230 Py_INCREF(newvalue);
231 } else {
232 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
233 newvalue, NULL);
234 if (newkey == NULL) {
235 Py_DECREF(newvalue);
236 return NULL;
237 }
238 }
239
240 assert(gbo->currkey == NULL);
241 gbo->currkey = newkey;
242 gbo->currvalue = newvalue;
243 }
244
245 assert(gbo->currkey != NULL);
246 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
247 if (rcmp <= 0)
248 /* got any error or current group is end */
249 return NULL;
250
251 r = gbo->currvalue;
252 gbo->currvalue = NULL;
253 Py_DECREF(gbo->currkey);
254 gbo->currkey = NULL;
255
256 return r;
257}
258
259static PyTypeObject _grouper_type = {
260 PyObject_HEAD_INIT(NULL)
261 0, /* ob_size */
262 "itertools._grouper", /* tp_name */
263 sizeof(_grouperobject), /* tp_basicsize */
264 0, /* tp_itemsize */
265 /* methods */
266 (destructor)_grouper_dealloc, /* tp_dealloc */
267 0, /* tp_print */
268 0, /* tp_getattr */
269 0, /* tp_setattr */
270 0, /* tp_compare */
271 0, /* tp_repr */
272 0, /* tp_as_number */
273 0, /* tp_as_sequence */
274 0, /* tp_as_mapping */
275 0, /* tp_hash */
276 0, /* tp_call */
277 0, /* tp_str */
278 PyObject_GenericGetAttr, /* tp_getattro */
279 0, /* tp_setattro */
280 0, /* tp_as_buffer */
281 Py_TPFLAGS_DEFAULT, /* tp_flags */
282 0, /* tp_doc */
283 0, /* tp_traverse */
284 0, /* tp_clear */
285 0, /* tp_richcompare */
286 0, /* tp_weaklistoffset */
287 PyObject_SelfIter, /* tp_iter */
288 (iternextfunc)_grouper_next, /* tp_iternext */
289 0, /* tp_methods */
290 0, /* tp_members */
291 0, /* tp_getset */
292 0, /* tp_base */
293 0, /* tp_dict */
294 0, /* tp_descr_get */
295 0, /* tp_descr_set */
296 0, /* tp_dictoffset */
297 0, /* tp_init */
298 0, /* tp_alloc */
299 0, /* tp_new */
300 PyObject_Del, /* tp_free */
301};
302
303
304
Raymond Hettingerad983e72003-11-12 14:32:26 +0000305/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000306
Raymond Hettingerad983e72003-11-12 14:32:26 +0000307/* The teedataobject pre-allocates space for LINKCELLS number of objects.
308 To help the object fit neatly inside cache lines (space for 16 to 32
309 pointers), the value should be a multiple of 16 minus space for
310 the other structure members including PyHEAD overhead. The larger the
311 value, the less memory overhead per object and the less time spent
312 allocating/deallocating new links. The smaller the number, the less
313 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000314*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000315#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000316
317typedef struct {
318 PyObject_HEAD
319 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000320 int numread;
321 PyObject *nextlink;
322 PyObject *(values[LINKCELLS]);
323} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000324
325typedef struct {
326 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327 teedataobject *dataobj;
328 int index;
329} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000330
Raymond Hettingerad983e72003-11-12 14:32:26 +0000331static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000332
333static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000336 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000337
Raymond Hettingerad983e72003-11-12 14:32:26 +0000338 tdo = PyObject_New(teedataobject, &teedataobject_type);
339 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000340 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000341
342 tdo->numread = 0;
343 tdo->nextlink = NULL;
344 Py_INCREF(it);
345 tdo->it = it;
346 return (PyObject *)tdo;
347}
348
349static PyObject *
350teedataobject_jumplink(teedataobject *tdo)
351{
352 if (tdo->nextlink == NULL)
353 tdo->nextlink = teedataobject_new(tdo->it);
354 Py_INCREF(tdo->nextlink);
355 return tdo->nextlink;
356}
357
358static PyObject *
359teedataobject_getitem(teedataobject *tdo, int i)
360{
361 PyObject *value;
362
363 assert(i < LINKCELLS);
364 if (i < tdo->numread)
365 value = tdo->values[i];
366 else {
367 /* this is the lead iterator, so fetch more data */
368 assert(i == tdo->numread);
369 value = PyIter_Next(tdo->it);
370 if (value == NULL)
371 return NULL;
372 tdo->numread++;
373 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000374 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000375 Py_INCREF(value);
376 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000377}
378
379static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000380teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000381{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000382 int i;
Raymond Hettinger45143692003-10-25 06:37:47 +0000383
Raymond Hettingerad983e72003-11-12 14:32:26 +0000384 for (i=0 ; i<tdo->numread ; i++)
385 Py_DECREF(tdo->values[i]);
386 Py_XDECREF(tdo->it);
387 Py_XDECREF(tdo->nextlink);
388 PyObject_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000389}
390
Raymond Hettingerad983e72003-11-12 14:32:26 +0000391PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000392
Raymond Hettingerad983e72003-11-12 14:32:26 +0000393static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000394 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000395 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000396 "itertools.tee_dataobject", /* tp_name */
397 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398 0, /* tp_itemsize */
399 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000400 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000401 0, /* tp_print */
402 0, /* tp_getattr */
403 0, /* tp_setattr */
404 0, /* tp_compare */
405 0, /* tp_repr */
406 0, /* tp_as_number */
407 0, /* tp_as_sequence */
408 0, /* tp_as_mapping */
409 0, /* tp_hash */
410 0, /* tp_call */
411 0, /* tp_str */
412 PyObject_GenericGetAttr, /* tp_getattro */
413 0, /* tp_setattro */
414 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000415 Py_TPFLAGS_DEFAULT, /* tp_flags */
416 teedataobject_doc, /* tp_doc */
417 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000418};
419
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420
421static PyTypeObject tee_type;
422
423static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000424tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000425{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000426 PyObject *value, *link;
427
428 if (to->index >= LINKCELLS) {
429 link = teedataobject_jumplink(to->dataobj);
430 Py_XDECREF(to->dataobj);
431 to->dataobj = (teedataobject *)link;
432 to->index = 0;
433 }
434 value = teedataobject_getitem(to->dataobj, to->index);
435 if (value == NULL)
436 return NULL;
437 to->index++;
438 return value;
439}
440
441static PyObject *
442tee_copy(teeobject *to)
443{
444 teeobject *newto;
445
446 newto = PyObject_New(teeobject, &tee_type);
447 if (newto == NULL)
448 return NULL;
449 Py_INCREF(to->dataobj);
450 newto->dataobj = to->dataobj;
451 newto->index = to->index;
452 return (PyObject *)newto;
453}
454
455PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
456
457static PyObject *
458tee_fromiterable(PyObject *iterable)
459{
460 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000461 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000462
463 it = PyObject_GetIter(iterable);
464 if (it == NULL)
465 return NULL;
466 if (PyObject_TypeCheck(it, &tee_type)) {
467 to = (teeobject *)tee_copy((teeobject *)it);
468 goto done;
469 }
470
471 to = PyObject_New(teeobject, &tee_type);
472 if (to == NULL)
473 goto done;
474 to->dataobj = (teedataobject *)teedataobject_new(it);
475 to->index = 0;
476done:
477 Py_XDECREF(it);
478 return (PyObject *)to;
479}
480
481static PyObject *
482tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
483{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000484 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000485
486 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
487 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000488 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000489}
490
491static void
492tee_dealloc(teeobject *to)
493{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000494 Py_XDECREF(to->dataobj);
495 PyObject_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000496}
497
Raymond Hettingerad983e72003-11-12 14:32:26 +0000498PyDoc_STRVAR(teeobject_doc,
499"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000500
Raymond Hettingerad983e72003-11-12 14:32:26 +0000501static PyMethodDef tee_methods[] = {
502 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
503 {NULL, NULL} /* sentinel */
504};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000505
506static PyTypeObject tee_type = {
507 PyObject_HEAD_INIT(NULL)
508 0, /* ob_size */
509 "itertools.tee", /* tp_name */
510 sizeof(teeobject), /* tp_basicsize */
511 0, /* tp_itemsize */
512 /* methods */
513 (destructor)tee_dealloc, /* tp_dealloc */
514 0, /* tp_print */
515 0, /* tp_getattr */
516 0, /* tp_setattr */
517 0, /* tp_compare */
518 0, /* tp_repr */
519 0, /* tp_as_number */
520 0, /* tp_as_sequence */
521 0, /* tp_as_mapping */
522 0, /* tp_hash */
523 0, /* tp_call */
524 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000525 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000526 0, /* tp_setattro */
527 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000528 Py_TPFLAGS_DEFAULT, /* tp_flags */
529 teeobject_doc, /* tp_doc */
530 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000531 0, /* tp_clear */
532 0, /* tp_richcompare */
533 0, /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000534 PyObject_SelfIter, /* tp_iter */
535 (iternextfunc)tee_next, /* tp_iternext */
536 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000537 0, /* tp_members */
538 0, /* tp_getset */
539 0, /* tp_base */
540 0, /* tp_dict */
541 0, /* tp_descr_get */
542 0, /* tp_descr_set */
543 0, /* tp_dictoffset */
544 0, /* tp_init */
545 0, /* tp_alloc */
546 tee_new, /* tp_new */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000547 PyObject_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000548};
549
Raymond Hettingerad983e72003-11-12 14:32:26 +0000550static PyObject *
551tee(PyObject *self, PyObject *args)
552{
553 int i, n=2;
554 PyObject *it, *iterable, *copyable, *result;
555
556 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
557 return NULL;
558 result = PyTuple_New(n);
559 if (result == NULL)
560 return NULL;
561 if (n == 0)
562 return result;
563 it = PyObject_GetIter(iterable);
564 if (it == NULL) {
565 Py_DECREF(result);
566 return NULL;
567 }
568 if (!PyObject_HasAttrString(it, "__copy__")) {
569 copyable = tee_fromiterable(it);
570 Py_DECREF(it);
571 if (copyable == NULL) {
572 Py_DECREF(result);
573 return NULL;
574 }
575 } else
576 copyable = it;
577 PyTuple_SET_ITEM(result, 0, copyable);
578 for (i=1 ; i<n ; i++) {
579 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
580 if (copyable == NULL) {
581 Py_DECREF(result);
582 return NULL;
583 }
584 PyTuple_SET_ITEM(result, i, copyable);
585 }
586 return result;
587}
588
589PyDoc_STRVAR(tee_doc,
590"tee(iterable, n=2) --> tuple of n independent iterators.");
591
592
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000593/* cycle object **********************************************************/
594
595typedef struct {
596 PyObject_HEAD
597 PyObject *it;
598 PyObject *saved;
599 int firstpass;
600} cycleobject;
601
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000602static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000603
604static PyObject *
605cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
606{
607 PyObject *it;
608 PyObject *iterable;
609 PyObject *saved;
610 cycleobject *lz;
611
612 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
613 return NULL;
614
615 /* Get iterator. */
616 it = PyObject_GetIter(iterable);
617 if (it == NULL)
618 return NULL;
619
620 saved = PyList_New(0);
621 if (saved == NULL) {
622 Py_DECREF(it);
623 return NULL;
624 }
625
626 /* create cycleobject structure */
627 lz = (cycleobject *)type->tp_alloc(type, 0);
628 if (lz == NULL) {
629 Py_DECREF(it);
630 Py_DECREF(saved);
631 return NULL;
632 }
633 lz->it = it;
634 lz->saved = saved;
635 lz->firstpass = 0;
636
637 return (PyObject *)lz;
638}
639
640static void
641cycle_dealloc(cycleobject *lz)
642{
643 PyObject_GC_UnTrack(lz);
644 Py_XDECREF(lz->saved);
645 Py_XDECREF(lz->it);
646 lz->ob_type->tp_free(lz);
647}
648
649static int
650cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
651{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000652 Py_VISIT(lz->it);
653 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000654 return 0;
655}
656
657static PyObject *
658cycle_next(cycleobject *lz)
659{
660 PyObject *item;
661 PyObject *it;
662
663 while (1) {
664 item = PyIter_Next(lz->it);
665 if (item != NULL) {
666 if (!lz->firstpass)
667 PyList_Append(lz->saved, item);
668 return item;
669 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000670 if (PyErr_Occurred()) {
671 if (PyErr_ExceptionMatches(PyExc_StopIteration))
672 PyErr_Clear();
673 else
674 return NULL;
675 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000676 if (PyList_Size(lz->saved) == 0)
677 return NULL;
678 it = PyObject_GetIter(lz->saved);
679 if (it == NULL)
680 return NULL;
681 Py_DECREF(lz->it);
682 lz->it = it;
683 lz->firstpass = 1;
684 }
685}
686
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000687PyDoc_STRVAR(cycle_doc,
688"cycle(iterable) --> cycle object\n\
689\n\
690Return elements from the iterable until it is exhausted.\n\
691Then repeat the sequence indefinitely.");
692
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000693static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000694 PyObject_HEAD_INIT(NULL)
695 0, /* ob_size */
696 "itertools.cycle", /* tp_name */
697 sizeof(cycleobject), /* tp_basicsize */
698 0, /* tp_itemsize */
699 /* methods */
700 (destructor)cycle_dealloc, /* tp_dealloc */
701 0, /* tp_print */
702 0, /* tp_getattr */
703 0, /* tp_setattr */
704 0, /* tp_compare */
705 0, /* tp_repr */
706 0, /* tp_as_number */
707 0, /* tp_as_sequence */
708 0, /* tp_as_mapping */
709 0, /* tp_hash */
710 0, /* tp_call */
711 0, /* tp_str */
712 PyObject_GenericGetAttr, /* tp_getattro */
713 0, /* tp_setattro */
714 0, /* tp_as_buffer */
715 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
716 Py_TPFLAGS_BASETYPE, /* tp_flags */
717 cycle_doc, /* tp_doc */
718 (traverseproc)cycle_traverse, /* tp_traverse */
719 0, /* tp_clear */
720 0, /* tp_richcompare */
721 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000722 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000723 (iternextfunc)cycle_next, /* tp_iternext */
724 0, /* tp_methods */
725 0, /* tp_members */
726 0, /* tp_getset */
727 0, /* tp_base */
728 0, /* tp_dict */
729 0, /* tp_descr_get */
730 0, /* tp_descr_set */
731 0, /* tp_dictoffset */
732 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000733 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000734 cycle_new, /* tp_new */
735 PyObject_GC_Del, /* tp_free */
736};
737
738
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000739/* dropwhile object **********************************************************/
740
741typedef struct {
742 PyObject_HEAD
743 PyObject *func;
744 PyObject *it;
745 long start;
746} dropwhileobject;
747
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000748static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000749
750static PyObject *
751dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
752{
753 PyObject *func, *seq;
754 PyObject *it;
755 dropwhileobject *lz;
756
757 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
758 return NULL;
759
760 /* Get iterator. */
761 it = PyObject_GetIter(seq);
762 if (it == NULL)
763 return NULL;
764
765 /* create dropwhileobject structure */
766 lz = (dropwhileobject *)type->tp_alloc(type, 0);
767 if (lz == NULL) {
768 Py_DECREF(it);
769 return NULL;
770 }
771 Py_INCREF(func);
772 lz->func = func;
773 lz->it = it;
774 lz->start = 0;
775
776 return (PyObject *)lz;
777}
778
779static void
780dropwhile_dealloc(dropwhileobject *lz)
781{
782 PyObject_GC_UnTrack(lz);
783 Py_XDECREF(lz->func);
784 Py_XDECREF(lz->it);
785 lz->ob_type->tp_free(lz);
786}
787
788static int
789dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
790{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000791 Py_VISIT(lz->it);
792 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000793 return 0;
794}
795
796static PyObject *
797dropwhile_next(dropwhileobject *lz)
798{
799 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000800 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000801 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000802 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000803
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000804 assert(PyIter_Check(it));
805 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000806 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000807 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000808 if (item == NULL)
809 return NULL;
810 if (lz->start == 1)
811 return item;
812
813 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
814 if (good == NULL) {
815 Py_DECREF(item);
816 return NULL;
817 }
818 ok = PyObject_IsTrue(good);
819 Py_DECREF(good);
820 if (!ok) {
821 lz->start = 1;
822 return item;
823 }
824 Py_DECREF(item);
825 }
826}
827
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000828PyDoc_STRVAR(dropwhile_doc,
829"dropwhile(predicate, iterable) --> dropwhile object\n\
830\n\
831Drop items from the iterable while predicate(item) is true.\n\
832Afterwards, return every element until the iterable is exhausted.");
833
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000834static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000835 PyObject_HEAD_INIT(NULL)
836 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000837 "itertools.dropwhile", /* tp_name */
838 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000839 0, /* tp_itemsize */
840 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000841 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 0, /* tp_print */
843 0, /* tp_getattr */
844 0, /* tp_setattr */
845 0, /* tp_compare */
846 0, /* tp_repr */
847 0, /* tp_as_number */
848 0, /* tp_as_sequence */
849 0, /* tp_as_mapping */
850 0, /* tp_hash */
851 0, /* tp_call */
852 0, /* tp_str */
853 PyObject_GenericGetAttr, /* tp_getattro */
854 0, /* tp_setattro */
855 0, /* tp_as_buffer */
856 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
857 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000858 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000859 (traverseproc)dropwhile_traverse, /* tp_traverse */
860 0, /* tp_clear */
861 0, /* tp_richcompare */
862 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000863 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000864 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000865 0, /* tp_methods */
866 0, /* tp_members */
867 0, /* tp_getset */
868 0, /* tp_base */
869 0, /* tp_dict */
870 0, /* tp_descr_get */
871 0, /* tp_descr_set */
872 0, /* tp_dictoffset */
873 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000874 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000875 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876 PyObject_GC_Del, /* tp_free */
877};
878
879
880/* takewhile object **********************************************************/
881
882typedef struct {
883 PyObject_HEAD
884 PyObject *func;
885 PyObject *it;
886 long stop;
887} takewhileobject;
888
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000889static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000890
891static PyObject *
892takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
893{
894 PyObject *func, *seq;
895 PyObject *it;
896 takewhileobject *lz;
897
898 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
899 return NULL;
900
901 /* Get iterator. */
902 it = PyObject_GetIter(seq);
903 if (it == NULL)
904 return NULL;
905
906 /* create takewhileobject structure */
907 lz = (takewhileobject *)type->tp_alloc(type, 0);
908 if (lz == NULL) {
909 Py_DECREF(it);
910 return NULL;
911 }
912 Py_INCREF(func);
913 lz->func = func;
914 lz->it = it;
915 lz->stop = 0;
916
917 return (PyObject *)lz;
918}
919
920static void
921takewhile_dealloc(takewhileobject *lz)
922{
923 PyObject_GC_UnTrack(lz);
924 Py_XDECREF(lz->func);
925 Py_XDECREF(lz->it);
926 lz->ob_type->tp_free(lz);
927}
928
929static int
930takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
931{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000932 Py_VISIT(lz->it);
933 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000934 return 0;
935}
936
937static PyObject *
938takewhile_next(takewhileobject *lz)
939{
940 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000941 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000942 long ok;
943
944 if (lz->stop == 1)
945 return NULL;
946
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000947 assert(PyIter_Check(it));
948 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000949 if (item == NULL)
950 return NULL;
951
952 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
953 if (good == NULL) {
954 Py_DECREF(item);
955 return NULL;
956 }
957 ok = PyObject_IsTrue(good);
958 Py_DECREF(good);
959 if (ok)
960 return item;
961 Py_DECREF(item);
962 lz->stop = 1;
963 return NULL;
964}
965
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000966PyDoc_STRVAR(takewhile_doc,
967"takewhile(predicate, iterable) --> takewhile object\n\
968\n\
969Return successive entries from an iterable as long as the \n\
970predicate evaluates to true for each entry.");
971
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000972static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000973 PyObject_HEAD_INIT(NULL)
974 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000975 "itertools.takewhile", /* tp_name */
976 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000977 0, /* tp_itemsize */
978 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000979 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000980 0, /* tp_print */
981 0, /* tp_getattr */
982 0, /* tp_setattr */
983 0, /* tp_compare */
984 0, /* tp_repr */
985 0, /* tp_as_number */
986 0, /* tp_as_sequence */
987 0, /* tp_as_mapping */
988 0, /* tp_hash */
989 0, /* tp_call */
990 0, /* tp_str */
991 PyObject_GenericGetAttr, /* tp_getattro */
992 0, /* tp_setattro */
993 0, /* tp_as_buffer */
994 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
995 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000996 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000997 (traverseproc)takewhile_traverse, /* tp_traverse */
998 0, /* tp_clear */
999 0, /* tp_richcompare */
1000 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001001 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001002 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001003 0, /* tp_methods */
1004 0, /* tp_members */
1005 0, /* tp_getset */
1006 0, /* tp_base */
1007 0, /* tp_dict */
1008 0, /* tp_descr_get */
1009 0, /* tp_descr_set */
1010 0, /* tp_dictoffset */
1011 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001012 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001013 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014 PyObject_GC_Del, /* tp_free */
1015};
1016
1017
1018/* islice object ************************************************************/
1019
1020typedef struct {
1021 PyObject_HEAD
1022 PyObject *it;
1023 long next;
1024 long stop;
1025 long step;
1026 long cnt;
1027} isliceobject;
1028
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001029static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001030
1031static PyObject *
1032islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1033{
1034 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001035 long start=0, stop=-1, step=1;
1036 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001037 int numargs;
1038 isliceobject *lz;
1039
1040 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +00001041 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001042 return NULL;
1043
1044 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001045 if (a1 != Py_None) {
1046 stop = PyInt_AsLong(a1);
1047 if (stop == -1) {
1048 if (PyErr_Occurred())
1049 PyErr_Clear();
1050 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001051 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001052 return NULL;
1053 }
1054 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001055 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001056 start = PyInt_AsLong(a1);
1057 if (start == -1 && PyErr_Occurred()) {
1058 PyErr_Clear();
1059 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001060 "Start argument must be a non-negative integer.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001061 return NULL;
1062 }
1063 if (a2 != Py_None) {
1064 stop = PyInt_AsLong(a2);
1065 if (stop == -1) {
1066 if (PyErr_Occurred())
1067 PyErr_Clear();
1068 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001069 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001070 return NULL;
1071 }
1072 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001073 }
1074
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001075 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001076 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001077 "Indices for islice() must be non-negative integers.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001078 return NULL;
1079 }
1080
1081 if (step<1) {
1082 PyErr_SetString(PyExc_ValueError,
1083 "Step must be one or larger for islice().");
1084 return NULL;
1085 }
1086
1087 /* Get iterator. */
1088 it = PyObject_GetIter(seq);
1089 if (it == NULL)
1090 return NULL;
1091
1092 /* create isliceobject structure */
1093 lz = (isliceobject *)type->tp_alloc(type, 0);
1094 if (lz == NULL) {
1095 Py_DECREF(it);
1096 return NULL;
1097 }
1098 lz->it = it;
1099 lz->next = start;
1100 lz->stop = stop;
1101 lz->step = step;
1102 lz->cnt = 0L;
1103
1104 return (PyObject *)lz;
1105}
1106
1107static void
1108islice_dealloc(isliceobject *lz)
1109{
1110 PyObject_GC_UnTrack(lz);
1111 Py_XDECREF(lz->it);
1112 lz->ob_type->tp_free(lz);
1113}
1114
1115static int
1116islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1117{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001118 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119 return 0;
1120}
1121
1122static PyObject *
1123islice_next(isliceobject *lz)
1124{
1125 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001126 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001127 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001128 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001129
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001130 assert(PyIter_Check(it));
1131 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001132 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001133 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001134 if (item == NULL)
1135 return NULL;
1136 Py_DECREF(item);
1137 lz->cnt++;
1138 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001139 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001140 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001141 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001142 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001143 if (item == NULL)
1144 return NULL;
1145 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001146 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001147 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001148 if (lz->next < oldnext) /* Check for overflow */
1149 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001150 return item;
1151}
1152
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153PyDoc_STRVAR(islice_doc,
1154"islice(iterable, [start,] stop [, step]) --> islice object\n\
1155\n\
1156Return an iterator whose next() method returns selected values from an\n\
1157iterable. If start is specified, will skip all preceding elements;\n\
1158otherwise, start defaults to zero. Step defaults to one. If\n\
1159specified as another value, step determines how many values are \n\
1160skipped between successive calls. Works like a slice() on a list\n\
1161but returns an iterator.");
1162
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001163static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001164 PyObject_HEAD_INIT(NULL)
1165 0, /* ob_size */
1166 "itertools.islice", /* tp_name */
1167 sizeof(isliceobject), /* tp_basicsize */
1168 0, /* tp_itemsize */
1169 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001170 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001171 0, /* tp_print */
1172 0, /* tp_getattr */
1173 0, /* tp_setattr */
1174 0, /* tp_compare */
1175 0, /* tp_repr */
1176 0, /* tp_as_number */
1177 0, /* tp_as_sequence */
1178 0, /* tp_as_mapping */
1179 0, /* tp_hash */
1180 0, /* tp_call */
1181 0, /* tp_str */
1182 PyObject_GenericGetAttr, /* tp_getattro */
1183 0, /* tp_setattro */
1184 0, /* tp_as_buffer */
1185 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1186 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001187 islice_doc, /* tp_doc */
1188 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001189 0, /* tp_clear */
1190 0, /* tp_richcompare */
1191 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001192 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001193 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194 0, /* tp_methods */
1195 0, /* tp_members */
1196 0, /* tp_getset */
1197 0, /* tp_base */
1198 0, /* tp_dict */
1199 0, /* tp_descr_get */
1200 0, /* tp_descr_set */
1201 0, /* tp_dictoffset */
1202 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001203 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204 islice_new, /* tp_new */
1205 PyObject_GC_Del, /* tp_free */
1206};
1207
1208
1209/* starmap object ************************************************************/
1210
1211typedef struct {
1212 PyObject_HEAD
1213 PyObject *func;
1214 PyObject *it;
1215} starmapobject;
1216
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001217static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218
1219static PyObject *
1220starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1221{
1222 PyObject *func, *seq;
1223 PyObject *it;
1224 starmapobject *lz;
1225
1226 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1227 return NULL;
1228
1229 /* Get iterator. */
1230 it = PyObject_GetIter(seq);
1231 if (it == NULL)
1232 return NULL;
1233
1234 /* create starmapobject structure */
1235 lz = (starmapobject *)type->tp_alloc(type, 0);
1236 if (lz == NULL) {
1237 Py_DECREF(it);
1238 return NULL;
1239 }
1240 Py_INCREF(func);
1241 lz->func = func;
1242 lz->it = it;
1243
1244 return (PyObject *)lz;
1245}
1246
1247static void
1248starmap_dealloc(starmapobject *lz)
1249{
1250 PyObject_GC_UnTrack(lz);
1251 Py_XDECREF(lz->func);
1252 Py_XDECREF(lz->it);
1253 lz->ob_type->tp_free(lz);
1254}
1255
1256static int
1257starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1258{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001259 Py_VISIT(lz->it);
1260 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261 return 0;
1262}
1263
1264static PyObject *
1265starmap_next(starmapobject *lz)
1266{
1267 PyObject *args;
1268 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001269 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001270
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001271 assert(PyIter_Check(it));
1272 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001273 if (args == NULL)
1274 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001275 if (!PyTuple_CheckExact(args)) {
1276 Py_DECREF(args);
1277 PyErr_SetString(PyExc_TypeError,
1278 "iterator must return a tuple");
1279 return NULL;
1280 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001281 result = PyObject_Call(lz->func, args, NULL);
1282 Py_DECREF(args);
1283 return result;
1284}
1285
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286PyDoc_STRVAR(starmap_doc,
1287"starmap(function, sequence) --> starmap object\n\
1288\n\
1289Return an iterator whose values are returned from the function evaluated\n\
1290with a argument tuple taken from the given sequence.");
1291
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001292static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001293 PyObject_HEAD_INIT(NULL)
1294 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001295 "itertools.starmap", /* tp_name */
1296 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001297 0, /* tp_itemsize */
1298 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001299 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300 0, /* tp_print */
1301 0, /* tp_getattr */
1302 0, /* tp_setattr */
1303 0, /* tp_compare */
1304 0, /* tp_repr */
1305 0, /* tp_as_number */
1306 0, /* tp_as_sequence */
1307 0, /* tp_as_mapping */
1308 0, /* tp_hash */
1309 0, /* tp_call */
1310 0, /* tp_str */
1311 PyObject_GenericGetAttr, /* tp_getattro */
1312 0, /* tp_setattro */
1313 0, /* tp_as_buffer */
1314 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1315 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001316 starmap_doc, /* tp_doc */
1317 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318 0, /* tp_clear */
1319 0, /* tp_richcompare */
1320 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001321 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001322 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001323 0, /* tp_methods */
1324 0, /* tp_members */
1325 0, /* tp_getset */
1326 0, /* tp_base */
1327 0, /* tp_dict */
1328 0, /* tp_descr_get */
1329 0, /* tp_descr_set */
1330 0, /* tp_dictoffset */
1331 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001332 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001333 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334 PyObject_GC_Del, /* tp_free */
1335};
1336
1337
1338/* imap object ************************************************************/
1339
1340typedef struct {
1341 PyObject_HEAD
1342 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001343 PyObject *func;
1344} imapobject;
1345
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001346static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001347
1348static PyObject *
1349imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1350{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001351 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352 imapobject *lz;
1353 int numargs, i;
1354
1355 numargs = PyTuple_Size(args);
1356 if (numargs < 2) {
1357 PyErr_SetString(PyExc_TypeError,
1358 "imap() must have at least two arguments.");
1359 return NULL;
1360 }
1361
1362 iters = PyTuple_New(numargs-1);
1363 if (iters == NULL)
1364 return NULL;
1365
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366 for (i=1 ; i<numargs ; i++) {
1367 /* Get iterator. */
1368 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1369 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370 Py_DECREF(iters);
1371 return NULL;
1372 }
1373 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374 }
1375
1376 /* create imapobject structure */
1377 lz = (imapobject *)type->tp_alloc(type, 0);
1378 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001379 Py_DECREF(iters);
1380 return NULL;
1381 }
1382 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383 func = PyTuple_GET_ITEM(args, 0);
1384 Py_INCREF(func);
1385 lz->func = func;
1386
1387 return (PyObject *)lz;
1388}
1389
1390static void
1391imap_dealloc(imapobject *lz)
1392{
1393 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394 Py_XDECREF(lz->iters);
1395 Py_XDECREF(lz->func);
1396 lz->ob_type->tp_free(lz);
1397}
1398
1399static int
1400imap_traverse(imapobject *lz, visitproc visit, void *arg)
1401{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001402 Py_VISIT(lz->iters);
1403 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404 return 0;
1405}
1406
Raymond Hettinger2012f172003-02-07 05:32:58 +00001407/*
1408imap() is an iterator version of __builtins__.map() except that it does
1409not have the None fill-in feature. That was intentionally left out for
1410the following reasons:
1411
1412 1) Itertools are designed to be easily combined and chained together.
1413 Having all tools stop with the shortest input is a unifying principle
1414 that makes it easier to combine finite iterators (supplying data) with
1415 infinite iterators like count() and repeat() (for supplying sequential
1416 or constant arguments to a function).
1417
1418 2) In typical use cases for combining itertools, having one finite data
1419 supplier run out before another is likely to be an error condition which
1420 should not pass silently by automatically supplying None.
1421
1422 3) The use cases for automatic None fill-in are rare -- not many functions
1423 do something useful when a parameter suddenly switches type and becomes
1424 None.
1425
1426 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001427 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001428
1429 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1430*/
1431
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432static PyObject *
1433imap_next(imapobject *lz)
1434{
1435 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001436 PyObject *argtuple;
1437 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438 int numargs, i;
1439
1440 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001441 argtuple = PyTuple_New(numargs);
1442 if (argtuple == NULL)
1443 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001444
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001445 for (i=0 ; i<numargs ; i++) {
1446 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1447 if (val == NULL) {
1448 Py_DECREF(argtuple);
1449 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001451 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001453 if (lz->func == Py_None)
1454 return argtuple;
1455 result = PyObject_Call(lz->func, argtuple, NULL);
1456 Py_DECREF(argtuple);
1457 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458}
1459
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460PyDoc_STRVAR(imap_doc,
1461"imap(func, *iterables) --> imap object\n\
1462\n\
1463Make an iterator that computes the function using arguments from\n\
1464each of the iterables. Like map() except that it returns\n\
1465an iterator instead of a list and that it stops when the shortest\n\
1466iterable is exhausted instead of filling in None for shorter\n\
1467iterables.");
1468
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001469static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470 PyObject_HEAD_INIT(NULL)
1471 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001472 "itertools.imap", /* tp_name */
1473 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001474 0, /* tp_itemsize */
1475 /* methods */
1476 (destructor)imap_dealloc, /* tp_dealloc */
1477 0, /* tp_print */
1478 0, /* tp_getattr */
1479 0, /* tp_setattr */
1480 0, /* tp_compare */
1481 0, /* tp_repr */
1482 0, /* tp_as_number */
1483 0, /* tp_as_sequence */
1484 0, /* tp_as_mapping */
1485 0, /* tp_hash */
1486 0, /* tp_call */
1487 0, /* tp_str */
1488 PyObject_GenericGetAttr, /* tp_getattro */
1489 0, /* tp_setattro */
1490 0, /* tp_as_buffer */
1491 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1492 Py_TPFLAGS_BASETYPE, /* tp_flags */
1493 imap_doc, /* tp_doc */
1494 (traverseproc)imap_traverse, /* tp_traverse */
1495 0, /* tp_clear */
1496 0, /* tp_richcompare */
1497 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001498 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001499 (iternextfunc)imap_next, /* tp_iternext */
1500 0, /* tp_methods */
1501 0, /* tp_members */
1502 0, /* tp_getset */
1503 0, /* tp_base */
1504 0, /* tp_dict */
1505 0, /* tp_descr_get */
1506 0, /* tp_descr_set */
1507 0, /* tp_dictoffset */
1508 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001509 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001510 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001511 PyObject_GC_Del, /* tp_free */
1512};
1513
1514
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001515/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516
1517typedef struct {
1518 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001519 long tuplesize;
1520 long iternum; /* which iterator is active */
1521 PyObject *ittuple; /* tuple of iterators */
1522} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001524static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525
1526static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001527chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001529 chainobject *lz;
1530 int tuplesize = PySequence_Length(args);
1531 int i;
1532 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001534 /* obtain iterators */
1535 assert(PyTuple_Check(args));
1536 ittuple = PyTuple_New(tuplesize);
1537 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001539 for (i=0; i < tuplesize; ++i) {
1540 PyObject *item = PyTuple_GET_ITEM(args, i);
1541 PyObject *it = PyObject_GetIter(item);
1542 if (it == NULL) {
1543 if (PyErr_ExceptionMatches(PyExc_TypeError))
1544 PyErr_Format(PyExc_TypeError,
1545 "chain argument #%d must support iteration",
1546 i+1);
1547 Py_DECREF(ittuple);
1548 return NULL;
1549 }
1550 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551 }
1552
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001553 /* create chainobject structure */
1554 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001555 if (lz == NULL) {
1556 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001557 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001558 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001559
1560 lz->ittuple = ittuple;
1561 lz->iternum = 0;
1562 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563
1564 return (PyObject *)lz;
1565}
1566
1567static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001568chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001569{
1570 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001571 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001572 lz->ob_type->tp_free(lz);
1573}
1574
Raymond Hettinger2012f172003-02-07 05:32:58 +00001575static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001576chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001577{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001578 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001579 return 0;
1580}
1581
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001582static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001583chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001584{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001585 PyObject *it;
1586 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001587
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001588 while (lz->iternum < lz->tuplesize) {
1589 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1590 item = PyIter_Next(it);
1591 if (item != NULL)
1592 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001593 if (PyErr_Occurred()) {
1594 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1595 PyErr_Clear();
1596 else
1597 return NULL;
1598 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001599 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001600 }
1601 return NULL;
1602}
1603
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001604PyDoc_STRVAR(chain_doc,
1605"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001607Return a chain object whose .next() method returns elements from the\n\
1608first iterable until it is exhausted, then elements from the next\n\
1609iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001610
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001611static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001612 PyObject_HEAD_INIT(NULL)
1613 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001614 "itertools.chain", /* tp_name */
1615 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616 0, /* tp_itemsize */
1617 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001618 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619 0, /* tp_print */
1620 0, /* tp_getattr */
1621 0, /* tp_setattr */
1622 0, /* tp_compare */
1623 0, /* tp_repr */
1624 0, /* tp_as_number */
1625 0, /* tp_as_sequence */
1626 0, /* tp_as_mapping */
1627 0, /* tp_hash */
1628 0, /* tp_call */
1629 0, /* tp_str */
1630 PyObject_GenericGetAttr, /* tp_getattro */
1631 0, /* tp_setattro */
1632 0, /* tp_as_buffer */
1633 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1634 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001635 chain_doc, /* tp_doc */
1636 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637 0, /* tp_clear */
1638 0, /* tp_richcompare */
1639 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001640 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001641 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001642 0, /* tp_methods */
1643 0, /* tp_members */
1644 0, /* tp_getset */
1645 0, /* tp_base */
1646 0, /* tp_dict */
1647 0, /* tp_descr_get */
1648 0, /* tp_descr_set */
1649 0, /* tp_dictoffset */
1650 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001651 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001652 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001653 PyObject_GC_Del, /* tp_free */
1654};
1655
1656
1657/* ifilter object ************************************************************/
1658
1659typedef struct {
1660 PyObject_HEAD
1661 PyObject *func;
1662 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001663} ifilterobject;
1664
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001665static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666
1667static PyObject *
1668ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1669{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001670 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671 PyObject *it;
1672 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001673
Raymond Hettinger60eca932003-02-09 06:40:58 +00001674 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675 return NULL;
1676
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677 /* Get iterator. */
1678 it = PyObject_GetIter(seq);
1679 if (it == NULL)
1680 return NULL;
1681
1682 /* create ifilterobject structure */
1683 lz = (ifilterobject *)type->tp_alloc(type, 0);
1684 if (lz == NULL) {
1685 Py_DECREF(it);
1686 return NULL;
1687 }
1688 Py_INCREF(func);
1689 lz->func = func;
1690 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001691
1692 return (PyObject *)lz;
1693}
1694
1695static void
1696ifilter_dealloc(ifilterobject *lz)
1697{
1698 PyObject_GC_UnTrack(lz);
1699 Py_XDECREF(lz->func);
1700 Py_XDECREF(lz->it);
1701 lz->ob_type->tp_free(lz);
1702}
1703
1704static int
1705ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1706{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001707 Py_VISIT(lz->it);
1708 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001709 return 0;
1710}
1711
1712static PyObject *
1713ifilter_next(ifilterobject *lz)
1714{
1715 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001716 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001717 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001718 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001719
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001720 assert(PyIter_Check(it));
1721 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001722 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001723 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724 if (item == NULL)
1725 return NULL;
1726
1727 if (lz->func == Py_None) {
1728 ok = PyObject_IsTrue(item);
1729 } else {
1730 PyObject *good;
1731 good = PyObject_CallFunctionObjArgs(lz->func,
1732 item, NULL);
1733 if (good == NULL) {
1734 Py_DECREF(item);
1735 return NULL;
1736 }
1737 ok = PyObject_IsTrue(good);
1738 Py_DECREF(good);
1739 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001740 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001741 return item;
1742 Py_DECREF(item);
1743 }
1744}
1745
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001747"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001748\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001749Return those items of sequence for which function(item) is true.\n\
1750If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001751
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001752static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753 PyObject_HEAD_INIT(NULL)
1754 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001755 "itertools.ifilter", /* tp_name */
1756 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757 0, /* tp_itemsize */
1758 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001759 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001760 0, /* tp_print */
1761 0, /* tp_getattr */
1762 0, /* tp_setattr */
1763 0, /* tp_compare */
1764 0, /* tp_repr */
1765 0, /* tp_as_number */
1766 0, /* tp_as_sequence */
1767 0, /* tp_as_mapping */
1768 0, /* tp_hash */
1769 0, /* tp_call */
1770 0, /* tp_str */
1771 PyObject_GenericGetAttr, /* tp_getattro */
1772 0, /* tp_setattro */
1773 0, /* tp_as_buffer */
1774 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1775 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001776 ifilter_doc, /* tp_doc */
1777 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001778 0, /* tp_clear */
1779 0, /* tp_richcompare */
1780 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001781 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001782 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001783 0, /* tp_methods */
1784 0, /* tp_members */
1785 0, /* tp_getset */
1786 0, /* tp_base */
1787 0, /* tp_dict */
1788 0, /* tp_descr_get */
1789 0, /* tp_descr_set */
1790 0, /* tp_dictoffset */
1791 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001792 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001793 ifilter_new, /* tp_new */
1794 PyObject_GC_Del, /* tp_free */
1795};
1796
1797
1798/* ifilterfalse object ************************************************************/
1799
1800typedef struct {
1801 PyObject_HEAD
1802 PyObject *func;
1803 PyObject *it;
1804} ifilterfalseobject;
1805
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001806static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001807
1808static PyObject *
1809ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1810{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001811 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001812 PyObject *it;
1813 ifilterfalseobject *lz;
1814
1815 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1816 return NULL;
1817
1818 /* Get iterator. */
1819 it = PyObject_GetIter(seq);
1820 if (it == NULL)
1821 return NULL;
1822
1823 /* create ifilterfalseobject structure */
1824 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1825 if (lz == NULL) {
1826 Py_DECREF(it);
1827 return NULL;
1828 }
1829 Py_INCREF(func);
1830 lz->func = func;
1831 lz->it = it;
1832
1833 return (PyObject *)lz;
1834}
1835
1836static void
1837ifilterfalse_dealloc(ifilterfalseobject *lz)
1838{
1839 PyObject_GC_UnTrack(lz);
1840 Py_XDECREF(lz->func);
1841 Py_XDECREF(lz->it);
1842 lz->ob_type->tp_free(lz);
1843}
1844
1845static int
1846ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1847{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001848 Py_VISIT(lz->it);
1849 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001850 return 0;
1851}
1852
1853static PyObject *
1854ifilterfalse_next(ifilterfalseobject *lz)
1855{
1856 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001857 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001858 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001859 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001860
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001861 assert(PyIter_Check(it));
1862 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001863 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001864 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001865 if (item == NULL)
1866 return NULL;
1867
1868 if (lz->func == Py_None) {
1869 ok = PyObject_IsTrue(item);
1870 } else {
1871 PyObject *good;
1872 good = PyObject_CallFunctionObjArgs(lz->func,
1873 item, NULL);
1874 if (good == NULL) {
1875 Py_DECREF(item);
1876 return NULL;
1877 }
1878 ok = PyObject_IsTrue(good);
1879 Py_DECREF(good);
1880 }
1881 if (!ok)
1882 return item;
1883 Py_DECREF(item);
1884 }
1885}
1886
Raymond Hettinger60eca932003-02-09 06:40:58 +00001887PyDoc_STRVAR(ifilterfalse_doc,
1888"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1889\n\
1890Return those items of sequence for which function(item) is false.\n\
1891If function is None, return the items that are false.");
1892
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001893static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001894 PyObject_HEAD_INIT(NULL)
1895 0, /* ob_size */
1896 "itertools.ifilterfalse", /* tp_name */
1897 sizeof(ifilterfalseobject), /* tp_basicsize */
1898 0, /* tp_itemsize */
1899 /* methods */
1900 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1901 0, /* tp_print */
1902 0, /* tp_getattr */
1903 0, /* tp_setattr */
1904 0, /* tp_compare */
1905 0, /* tp_repr */
1906 0, /* tp_as_number */
1907 0, /* tp_as_sequence */
1908 0, /* tp_as_mapping */
1909 0, /* tp_hash */
1910 0, /* tp_call */
1911 0, /* tp_str */
1912 PyObject_GenericGetAttr, /* tp_getattro */
1913 0, /* tp_setattro */
1914 0, /* tp_as_buffer */
1915 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1916 Py_TPFLAGS_BASETYPE, /* tp_flags */
1917 ifilterfalse_doc, /* tp_doc */
1918 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1919 0, /* tp_clear */
1920 0, /* tp_richcompare */
1921 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001922 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001923 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1924 0, /* tp_methods */
1925 0, /* tp_members */
1926 0, /* tp_getset */
1927 0, /* tp_base */
1928 0, /* tp_dict */
1929 0, /* tp_descr_get */
1930 0, /* tp_descr_set */
1931 0, /* tp_dictoffset */
1932 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001933 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001934 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001935 PyObject_GC_Del, /* tp_free */
1936};
1937
1938
1939/* count object ************************************************************/
1940
1941typedef struct {
1942 PyObject_HEAD
1943 long cnt;
1944} countobject;
1945
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001946static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001947
1948static PyObject *
1949count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1950{
1951 countobject *lz;
1952 long cnt = 0;
1953
1954 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1955 return NULL;
1956
1957 /* create countobject structure */
1958 lz = (countobject *)PyObject_New(countobject, &count_type);
1959 if (lz == NULL)
1960 return NULL;
1961 lz->cnt = cnt;
1962
1963 return (PyObject *)lz;
1964}
1965
1966static PyObject *
1967count_next(countobject *lz)
1968{
1969 return PyInt_FromLong(lz->cnt++);
1970}
1971
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001972static PyObject *
1973count_repr(countobject *lz)
1974{
Brett Cannon00468392004-04-13 02:43:53 +00001975 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001976}
1977
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001978PyDoc_STRVAR(count_doc,
1979"count([firstval]) --> count object\n\
1980\n\
1981Return a count object whose .next() method returns consecutive\n\
1982integers starting from zero or, if specified, from firstval.");
1983
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001984static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001985 PyObject_HEAD_INIT(NULL)
1986 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001987 "itertools.count", /* tp_name */
1988 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001989 0, /* tp_itemsize */
1990 /* methods */
1991 (destructor)PyObject_Del, /* tp_dealloc */
1992 0, /* tp_print */
1993 0, /* tp_getattr */
1994 0, /* tp_setattr */
1995 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001996 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001997 0, /* tp_as_number */
1998 0, /* tp_as_sequence */
1999 0, /* tp_as_mapping */
2000 0, /* tp_hash */
2001 0, /* tp_call */
2002 0, /* tp_str */
2003 PyObject_GenericGetAttr, /* tp_getattro */
2004 0, /* tp_setattro */
2005 0, /* tp_as_buffer */
2006 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002007 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002008 0, /* tp_traverse */
2009 0, /* tp_clear */
2010 0, /* tp_richcompare */
2011 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002012 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002013 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002014 0, /* tp_methods */
2015 0, /* tp_members */
2016 0, /* tp_getset */
2017 0, /* tp_base */
2018 0, /* tp_dict */
2019 0, /* tp_descr_get */
2020 0, /* tp_descr_set */
2021 0, /* tp_dictoffset */
2022 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002023 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002024 count_new, /* tp_new */
2025};
2026
2027
2028/* izip object ************************************************************/
2029
2030#include "Python.h"
2031
2032typedef struct {
2033 PyObject_HEAD
2034 long tuplesize;
2035 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002036 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002037} izipobject;
2038
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002039static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002040
2041static PyObject *
2042izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2043{
2044 izipobject *lz;
2045 int i;
2046 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002047 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002048 int tuplesize = PySequence_Length(args);
2049
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002050 /* args must be a tuple */
2051 assert(PyTuple_Check(args));
2052
2053 /* obtain iterators */
2054 ittuple = PyTuple_New(tuplesize);
2055 if(ittuple == NULL)
2056 return NULL;
2057 for (i=0; i < tuplesize; ++i) {
2058 PyObject *item = PyTuple_GET_ITEM(args, i);
2059 PyObject *it = PyObject_GetIter(item);
2060 if (it == NULL) {
2061 if (PyErr_ExceptionMatches(PyExc_TypeError))
2062 PyErr_Format(PyExc_TypeError,
2063 "izip argument #%d must support iteration",
2064 i+1);
2065 Py_DECREF(ittuple);
2066 return NULL;
2067 }
2068 PyTuple_SET_ITEM(ittuple, i, it);
2069 }
2070
Raymond Hettinger2012f172003-02-07 05:32:58 +00002071 /* create a result holder */
2072 result = PyTuple_New(tuplesize);
2073 if (result == NULL) {
2074 Py_DECREF(ittuple);
2075 return NULL;
2076 }
2077 for (i=0 ; i < tuplesize ; i++) {
2078 Py_INCREF(Py_None);
2079 PyTuple_SET_ITEM(result, i, Py_None);
2080 }
2081
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002082 /* create izipobject structure */
2083 lz = (izipobject *)type->tp_alloc(type, 0);
2084 if (lz == NULL) {
2085 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002086 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002087 return NULL;
2088 }
2089 lz->ittuple = ittuple;
2090 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002091 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002092
2093 return (PyObject *)lz;
2094}
2095
2096static void
2097izip_dealloc(izipobject *lz)
2098{
2099 PyObject_GC_UnTrack(lz);
2100 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002101 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002102 lz->ob_type->tp_free(lz);
2103}
2104
2105static int
2106izip_traverse(izipobject *lz, visitproc visit, void *arg)
2107{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002108 Py_VISIT(lz->ittuple);
2109 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002110 return 0;
2111}
2112
2113static PyObject *
2114izip_next(izipobject *lz)
2115{
2116 int i;
2117 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002118 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002119 PyObject *it;
2120 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002121 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002122
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002123 if (tuplesize == 0)
2124 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002125 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002126 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002127 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002128 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002129 assert(PyIter_Check(it));
2130 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002131 if (item == NULL) {
2132 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002133 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002134 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002135 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002136 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002137 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002138 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002139 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002140 result = PyTuple_New(tuplesize);
2141 if (result == NULL)
2142 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002143 for (i=0 ; i < tuplesize ; i++) {
2144 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002145 assert(PyIter_Check(it));
2146 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002147 if (item == NULL) {
2148 Py_DECREF(result);
2149 return NULL;
2150 }
2151 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002152 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002153 }
2154 return result;
2155}
2156
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002157PyDoc_STRVAR(izip_doc,
2158"izip(iter1 [,iter2 [...]]) --> izip object\n\
2159\n\
2160Return a izip object whose .next() method returns a tuple where\n\
2161the i-th element comes from the i-th iterable argument. The .next()\n\
2162method continues until the shortest iterable in the argument sequence\n\
2163is exhausted and then it raises StopIteration. Works like the zip()\n\
2164function but consumes less memory by returning an iterator instead of\n\
2165a list.");
2166
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002167static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002168 PyObject_HEAD_INIT(NULL)
2169 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002170 "itertools.izip", /* tp_name */
2171 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002172 0, /* tp_itemsize */
2173 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002174 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002175 0, /* tp_print */
2176 0, /* tp_getattr */
2177 0, /* tp_setattr */
2178 0, /* tp_compare */
2179 0, /* tp_repr */
2180 0, /* tp_as_number */
2181 0, /* tp_as_sequence */
2182 0, /* tp_as_mapping */
2183 0, /* tp_hash */
2184 0, /* tp_call */
2185 0, /* tp_str */
2186 PyObject_GenericGetAttr, /* tp_getattro */
2187 0, /* tp_setattro */
2188 0, /* tp_as_buffer */
2189 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2190 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002191 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002192 (traverseproc)izip_traverse, /* tp_traverse */
2193 0, /* tp_clear */
2194 0, /* tp_richcompare */
2195 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002196 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002197 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002198 0, /* tp_methods */
2199 0, /* tp_members */
2200 0, /* tp_getset */
2201 0, /* tp_base */
2202 0, /* tp_dict */
2203 0, /* tp_descr_get */
2204 0, /* tp_descr_set */
2205 0, /* tp_dictoffset */
2206 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002207 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002208 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002209 PyObject_GC_Del, /* tp_free */
2210};
2211
2212
2213/* repeat object ************************************************************/
2214
2215typedef struct {
2216 PyObject_HEAD
2217 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002218 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002219} repeatobject;
2220
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002221static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002222
2223static PyObject *
2224repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2225{
2226 repeatobject *ro;
2227 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002228 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002229
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002230 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002231 return NULL;
2232
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002233 if (PyTuple_Size(args) == 2 && cnt < 0)
2234 cnt = 0;
2235
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002236 ro = (repeatobject *)type->tp_alloc(type, 0);
2237 if (ro == NULL)
2238 return NULL;
2239 Py_INCREF(element);
2240 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002241 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002242 return (PyObject *)ro;
2243}
2244
2245static void
2246repeat_dealloc(repeatobject *ro)
2247{
2248 PyObject_GC_UnTrack(ro);
2249 Py_XDECREF(ro->element);
2250 ro->ob_type->tp_free(ro);
2251}
2252
2253static int
2254repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2255{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002256 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002257 return 0;
2258}
2259
2260static PyObject *
2261repeat_next(repeatobject *ro)
2262{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002263 if (ro->cnt == 0)
2264 return NULL;
2265 if (ro->cnt > 0)
2266 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002267 Py_INCREF(ro->element);
2268 return ro->element;
2269}
2270
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002271static PyObject *
2272repeat_repr(repeatobject *ro)
2273{
2274 PyObject *result, *objrepr;
2275
2276 objrepr = PyObject_Repr(ro->element);
2277 if (objrepr == NULL)
2278 return NULL;
2279
2280 if (ro->cnt == -1)
2281 result = PyString_FromFormat("repeat(%s)",
2282 PyString_AS_STRING(objrepr));
2283 else
Brett Cannon00468392004-04-13 02:43:53 +00002284 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002285 PyString_AS_STRING(objrepr), ro->cnt);
2286 Py_DECREF(objrepr);
2287 return result;
2288}
2289
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002290static int
2291repeat_len(repeatobject *ro)
2292{
2293 if (ro->cnt == -1)
2294 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2295 return (int)(ro->cnt);
2296}
2297
2298static PySequenceMethods repeat_as_sequence = {
2299 (inquiry)repeat_len, /* sq_length */
2300 0, /* sq_concat */
2301};
2302
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002303PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002304"repeat(element [,times]) -> create an iterator which returns the element\n\
2305for the specified number of times. If not specified, returns the element\n\
2306endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002307
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002308static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002309 PyObject_HEAD_INIT(NULL)
2310 0, /* ob_size */
2311 "itertools.repeat", /* tp_name */
2312 sizeof(repeatobject), /* tp_basicsize */
2313 0, /* tp_itemsize */
2314 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002315 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002316 0, /* tp_print */
2317 0, /* tp_getattr */
2318 0, /* tp_setattr */
2319 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002320 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002321 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002322 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002323 0, /* tp_as_mapping */
2324 0, /* tp_hash */
2325 0, /* tp_call */
2326 0, /* tp_str */
2327 PyObject_GenericGetAttr, /* tp_getattro */
2328 0, /* tp_setattro */
2329 0, /* tp_as_buffer */
2330 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2331 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002332 repeat_doc, /* tp_doc */
2333 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002334 0, /* tp_clear */
2335 0, /* tp_richcompare */
2336 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002337 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002338 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002339 0, /* tp_methods */
2340 0, /* tp_members */
2341 0, /* tp_getset */
2342 0, /* tp_base */
2343 0, /* tp_dict */
2344 0, /* tp_descr_get */
2345 0, /* tp_descr_set */
2346 0, /* tp_dictoffset */
2347 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002348 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002349 repeat_new, /* tp_new */
2350 PyObject_GC_Del, /* tp_free */
2351};
2352
2353
2354/* module level code ********************************************************/
2355
2356PyDoc_STRVAR(module_doc,
2357"Functional tools for creating and using iterators.\n\
2358\n\
2359Infinite iterators:\n\
2360count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002361cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002362repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002363\n\
2364Iterators terminating on the shortest input sequence:\n\
2365izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002366ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2367ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002368islice(seq, [start,] stop [, step]) --> elements from\n\
2369 seq[start:stop:step]\n\
2370imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2371starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002372tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002373chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002374takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2375dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002376groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002377");
2378
2379
Raymond Hettingerad983e72003-11-12 14:32:26 +00002380static PyMethodDef module_methods[] = {
2381 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2382 {NULL, NULL} /* sentinel */
2383};
2384
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002385PyMODINIT_FUNC
2386inititertools(void)
2387{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002388 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002389 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002390 char *name;
2391 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002392 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002393 &dropwhile_type,
2394 &takewhile_type,
2395 &islice_type,
2396 &starmap_type,
2397 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002398 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002399 &ifilter_type,
2400 &ifilterfalse_type,
2401 &count_type,
2402 &izip_type,
2403 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002404 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002405 NULL
2406 };
2407
Skip Montanarof3938fd2004-02-10 20:27:40 +00002408 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002409 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002410
Raymond Hettinger60eca932003-02-09 06:40:58 +00002411 for (i=0 ; typelist[i] != NULL ; i++) {
2412 if (PyType_Ready(typelist[i]) < 0)
2413 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002414 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002415 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002416 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002417 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002418 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002419
2420 if (PyType_Ready(&teedataobject_type) < 0)
2421 return;
2422 if (PyType_Ready(&tee_type) < 0)
2423 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002424 if (PyType_Ready(&_grouper_type) < 0)
2425 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002426}