blob: 4069ea2dafbeaeaafca89edb572ff1458f2ff06c [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{
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000078 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000079
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
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000113 tmp = gbo->currkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000114 gbo->currkey = newkey;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000115 Py_XDECREF(tmp);
116
117 tmp = gbo->currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000118 gbo->currvalue = newvalue;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000119 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000120 }
121
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122 Py_INCREF(gbo->currkey);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000123 tmp = gbo->tgtkey;
124 gbo->tgtkey = gbo->currkey;
125 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000126
127 grouper = _grouper_create(gbo, gbo->tgtkey);
128 if (grouper == NULL)
129 return NULL;
130
131 r = PyTuple_Pack(2, gbo->currkey, grouper);
132 Py_DECREF(grouper);
133 return r;
134}
135
136PyDoc_STRVAR(groupby_doc,
137"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
138(key, sub-iterator) grouped by each value of key(value).\n");
139
140static PyTypeObject groupby_type = {
141 PyObject_HEAD_INIT(NULL)
142 0, /* ob_size */
143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_compare */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
199 _grouperobject *igo;
200
201 igo = PyObject_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
208
209 return (PyObject *)igo;
210}
211
212static void
213_grouper_dealloc(_grouperobject *igo)
214{
215 Py_DECREF(igo->parent);
216 Py_DECREF(igo->tgtkey);
217 PyObject_Del(igo);
218}
219
220static PyObject *
221_grouper_next(_grouperobject *igo)
222{
223 groupbyobject *gbo = (groupbyobject *)igo->parent;
224 PyObject *newvalue, *newkey, *r;
225 int rcmp;
226
227 if (gbo->currvalue == NULL) {
228 newvalue = PyIter_Next(gbo->it);
229 if (newvalue == NULL)
230 return NULL;
231
232 if (gbo->keyfunc == Py_None) {
233 newkey = newvalue;
234 Py_INCREF(newvalue);
235 } else {
236 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
237 newvalue, NULL);
238 if (newkey == NULL) {
239 Py_DECREF(newvalue);
240 return NULL;
241 }
242 }
243
244 assert(gbo->currkey == NULL);
245 gbo->currkey = newkey;
246 gbo->currvalue = newvalue;
247 }
248
249 assert(gbo->currkey != NULL);
250 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
251 if (rcmp <= 0)
252 /* got any error or current group is end */
253 return NULL;
254
255 r = gbo->currvalue;
256 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000257 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
259 return r;
260}
261
262static PyTypeObject _grouper_type = {
263 PyObject_HEAD_INIT(NULL)
264 0, /* ob_size */
265 "itertools._grouper", /* tp_name */
266 sizeof(_grouperobject), /* tp_basicsize */
267 0, /* tp_itemsize */
268 /* methods */
269 (destructor)_grouper_dealloc, /* tp_dealloc */
270 0, /* tp_print */
271 0, /* tp_getattr */
272 0, /* tp_setattr */
273 0, /* tp_compare */
274 0, /* tp_repr */
275 0, /* tp_as_number */
276 0, /* tp_as_sequence */
277 0, /* tp_as_mapping */
278 0, /* tp_hash */
279 0, /* tp_call */
280 0, /* tp_str */
281 PyObject_GenericGetAttr, /* tp_getattro */
282 0, /* tp_setattro */
283 0, /* tp_as_buffer */
284 Py_TPFLAGS_DEFAULT, /* tp_flags */
285 0, /* tp_doc */
286 0, /* tp_traverse */
287 0, /* tp_clear */
288 0, /* tp_richcompare */
289 0, /* tp_weaklistoffset */
290 PyObject_SelfIter, /* tp_iter */
291 (iternextfunc)_grouper_next, /* tp_iternext */
292 0, /* tp_methods */
293 0, /* tp_members */
294 0, /* tp_getset */
295 0, /* tp_base */
296 0, /* tp_dict */
297 0, /* tp_descr_get */
298 0, /* tp_descr_set */
299 0, /* tp_dictoffset */
300 0, /* tp_init */
301 0, /* tp_alloc */
302 0, /* tp_new */
303 PyObject_Del, /* tp_free */
304};
305
306
307
Raymond Hettingerad983e72003-11-12 14:32:26 +0000308/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000309
Raymond Hettingerad983e72003-11-12 14:32:26 +0000310/* The teedataobject pre-allocates space for LINKCELLS number of objects.
311 To help the object fit neatly inside cache lines (space for 16 to 32
312 pointers), the value should be a multiple of 16 minus space for
313 the other structure members including PyHEAD overhead. The larger the
314 value, the less memory overhead per object and the less time spent
315 allocating/deallocating new links. The smaller the number, the less
316 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000317*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000318#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000319
320typedef struct {
321 PyObject_HEAD
322 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000323 int numread;
324 PyObject *nextlink;
325 PyObject *(values[LINKCELLS]);
326} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000327
328typedef struct {
329 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000330 teedataobject *dataobj;
331 int index;
332} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000333
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335
336static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000337teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000338{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000340
Raymond Hettingerad983e72003-11-12 14:32:26 +0000341 tdo = PyObject_New(teedataobject, &teedataobject_type);
342 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000343 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344
345 tdo->numread = 0;
346 tdo->nextlink = NULL;
347 Py_INCREF(it);
348 tdo->it = it;
349 return (PyObject *)tdo;
350}
351
352static PyObject *
353teedataobject_jumplink(teedataobject *tdo)
354{
355 if (tdo->nextlink == NULL)
356 tdo->nextlink = teedataobject_new(tdo->it);
357 Py_INCREF(tdo->nextlink);
358 return tdo->nextlink;
359}
360
361static PyObject *
362teedataobject_getitem(teedataobject *tdo, int i)
363{
364 PyObject *value;
365
366 assert(i < LINKCELLS);
367 if (i < tdo->numread)
368 value = tdo->values[i];
369 else {
370 /* this is the lead iterator, so fetch more data */
371 assert(i == tdo->numread);
372 value = PyIter_Next(tdo->it);
373 if (value == NULL)
374 return NULL;
375 tdo->numread++;
376 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000377 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000378 Py_INCREF(value);
379 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000380}
381
382static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000383teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000384{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000385 int i;
Raymond Hettinger45143692003-10-25 06:37:47 +0000386
Raymond Hettingerad983e72003-11-12 14:32:26 +0000387 for (i=0 ; i<tdo->numread ; i++)
388 Py_DECREF(tdo->values[i]);
389 Py_XDECREF(tdo->it);
390 Py_XDECREF(tdo->nextlink);
391 PyObject_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000392}
393
Raymond Hettingerad983e72003-11-12 14:32:26 +0000394PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000395
Raymond Hettingerad983e72003-11-12 14:32:26 +0000396static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000397 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000398 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399 "itertools.tee_dataobject", /* tp_name */
400 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000401 0, /* tp_itemsize */
402 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000403 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000404 0, /* tp_print */
405 0, /* tp_getattr */
406 0, /* tp_setattr */
407 0, /* tp_compare */
408 0, /* tp_repr */
409 0, /* tp_as_number */
410 0, /* tp_as_sequence */
411 0, /* tp_as_mapping */
412 0, /* tp_hash */
413 0, /* tp_call */
414 0, /* tp_str */
415 PyObject_GenericGetAttr, /* tp_getattro */
416 0, /* tp_setattro */
417 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000418 Py_TPFLAGS_DEFAULT, /* tp_flags */
419 teedataobject_doc, /* tp_doc */
420 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000421};
422
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000423
424static PyTypeObject tee_type;
425
426static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000427tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000428{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000429 PyObject *value, *link;
430
431 if (to->index >= LINKCELLS) {
432 link = teedataobject_jumplink(to->dataobj);
433 Py_XDECREF(to->dataobj);
434 to->dataobj = (teedataobject *)link;
435 to->index = 0;
436 }
437 value = teedataobject_getitem(to->dataobj, to->index);
438 if (value == NULL)
439 return NULL;
440 to->index++;
441 return value;
442}
443
444static PyObject *
445tee_copy(teeobject *to)
446{
447 teeobject *newto;
448
449 newto = PyObject_New(teeobject, &tee_type);
450 if (newto == NULL)
451 return NULL;
452 Py_INCREF(to->dataobj);
453 newto->dataobj = to->dataobj;
454 newto->index = to->index;
455 return (PyObject *)newto;
456}
457
458PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
459
460static PyObject *
461tee_fromiterable(PyObject *iterable)
462{
463 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000464 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000465
466 it = PyObject_GetIter(iterable);
467 if (it == NULL)
468 return NULL;
469 if (PyObject_TypeCheck(it, &tee_type)) {
470 to = (teeobject *)tee_copy((teeobject *)it);
471 goto done;
472 }
473
474 to = PyObject_New(teeobject, &tee_type);
475 if (to == NULL)
476 goto done;
477 to->dataobj = (teedataobject *)teedataobject_new(it);
478 to->index = 0;
479done:
480 Py_XDECREF(it);
481 return (PyObject *)to;
482}
483
484static PyObject *
485tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
486{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000487 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000488
489 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
490 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000491 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000492}
493
494static void
495tee_dealloc(teeobject *to)
496{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000497 Py_XDECREF(to->dataobj);
498 PyObject_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000499}
500
Raymond Hettingerad983e72003-11-12 14:32:26 +0000501PyDoc_STRVAR(teeobject_doc,
502"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000503
Raymond Hettingerad983e72003-11-12 14:32:26 +0000504static PyMethodDef tee_methods[] = {
505 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
506 {NULL, NULL} /* sentinel */
507};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000508
509static PyTypeObject tee_type = {
510 PyObject_HEAD_INIT(NULL)
511 0, /* ob_size */
512 "itertools.tee", /* tp_name */
513 sizeof(teeobject), /* tp_basicsize */
514 0, /* tp_itemsize */
515 /* methods */
516 (destructor)tee_dealloc, /* tp_dealloc */
517 0, /* tp_print */
518 0, /* tp_getattr */
519 0, /* tp_setattr */
520 0, /* tp_compare */
521 0, /* tp_repr */
522 0, /* tp_as_number */
523 0, /* tp_as_sequence */
524 0, /* tp_as_mapping */
525 0, /* tp_hash */
526 0, /* tp_call */
527 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000528 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000529 0, /* tp_setattro */
530 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000531 Py_TPFLAGS_DEFAULT, /* tp_flags */
532 teeobject_doc, /* tp_doc */
533 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000534 0, /* tp_clear */
535 0, /* tp_richcompare */
536 0, /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000537 PyObject_SelfIter, /* tp_iter */
538 (iternextfunc)tee_next, /* tp_iternext */
539 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000540 0, /* tp_members */
541 0, /* tp_getset */
542 0, /* tp_base */
543 0, /* tp_dict */
544 0, /* tp_descr_get */
545 0, /* tp_descr_set */
546 0, /* tp_dictoffset */
547 0, /* tp_init */
548 0, /* tp_alloc */
549 tee_new, /* tp_new */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000550 PyObject_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000551};
552
Raymond Hettingerad983e72003-11-12 14:32:26 +0000553static PyObject *
554tee(PyObject *self, PyObject *args)
555{
556 int i, n=2;
557 PyObject *it, *iterable, *copyable, *result;
558
559 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
560 return NULL;
561 result = PyTuple_New(n);
562 if (result == NULL)
563 return NULL;
564 if (n == 0)
565 return result;
566 it = PyObject_GetIter(iterable);
567 if (it == NULL) {
568 Py_DECREF(result);
569 return NULL;
570 }
571 if (!PyObject_HasAttrString(it, "__copy__")) {
572 copyable = tee_fromiterable(it);
573 Py_DECREF(it);
574 if (copyable == NULL) {
575 Py_DECREF(result);
576 return NULL;
577 }
578 } else
579 copyable = it;
580 PyTuple_SET_ITEM(result, 0, copyable);
581 for (i=1 ; i<n ; i++) {
582 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
583 if (copyable == NULL) {
584 Py_DECREF(result);
585 return NULL;
586 }
587 PyTuple_SET_ITEM(result, i, copyable);
588 }
589 return result;
590}
591
592PyDoc_STRVAR(tee_doc,
593"tee(iterable, n=2) --> tuple of n independent iterators.");
594
595
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000596/* cycle object **********************************************************/
597
598typedef struct {
599 PyObject_HEAD
600 PyObject *it;
601 PyObject *saved;
602 int firstpass;
603} cycleobject;
604
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000605static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000606
607static PyObject *
608cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
609{
610 PyObject *it;
611 PyObject *iterable;
612 PyObject *saved;
613 cycleobject *lz;
614
615 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
616 return NULL;
617
618 /* Get iterator. */
619 it = PyObject_GetIter(iterable);
620 if (it == NULL)
621 return NULL;
622
623 saved = PyList_New(0);
624 if (saved == NULL) {
625 Py_DECREF(it);
626 return NULL;
627 }
628
629 /* create cycleobject structure */
630 lz = (cycleobject *)type->tp_alloc(type, 0);
631 if (lz == NULL) {
632 Py_DECREF(it);
633 Py_DECREF(saved);
634 return NULL;
635 }
636 lz->it = it;
637 lz->saved = saved;
638 lz->firstpass = 0;
639
640 return (PyObject *)lz;
641}
642
643static void
644cycle_dealloc(cycleobject *lz)
645{
646 PyObject_GC_UnTrack(lz);
647 Py_XDECREF(lz->saved);
648 Py_XDECREF(lz->it);
649 lz->ob_type->tp_free(lz);
650}
651
652static int
653cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
654{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000655 Py_VISIT(lz->it);
656 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000657 return 0;
658}
659
660static PyObject *
661cycle_next(cycleobject *lz)
662{
663 PyObject *item;
664 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000665 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000666
667 while (1) {
668 item = PyIter_Next(lz->it);
669 if (item != NULL) {
670 if (!lz->firstpass)
671 PyList_Append(lz->saved, item);
672 return item;
673 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000674 if (PyErr_Occurred()) {
675 if (PyErr_ExceptionMatches(PyExc_StopIteration))
676 PyErr_Clear();
677 else
678 return NULL;
679 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000680 if (PyList_Size(lz->saved) == 0)
681 return NULL;
682 it = PyObject_GetIter(lz->saved);
683 if (it == NULL)
684 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000685 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000686 lz->it = it;
687 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000688 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000689 }
690}
691
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000692PyDoc_STRVAR(cycle_doc,
693"cycle(iterable) --> cycle object\n\
694\n\
695Return elements from the iterable until it is exhausted.\n\
696Then repeat the sequence indefinitely.");
697
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000698static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000699 PyObject_HEAD_INIT(NULL)
700 0, /* ob_size */
701 "itertools.cycle", /* tp_name */
702 sizeof(cycleobject), /* tp_basicsize */
703 0, /* tp_itemsize */
704 /* methods */
705 (destructor)cycle_dealloc, /* tp_dealloc */
706 0, /* tp_print */
707 0, /* tp_getattr */
708 0, /* tp_setattr */
709 0, /* tp_compare */
710 0, /* tp_repr */
711 0, /* tp_as_number */
712 0, /* tp_as_sequence */
713 0, /* tp_as_mapping */
714 0, /* tp_hash */
715 0, /* tp_call */
716 0, /* tp_str */
717 PyObject_GenericGetAttr, /* tp_getattro */
718 0, /* tp_setattro */
719 0, /* tp_as_buffer */
720 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
721 Py_TPFLAGS_BASETYPE, /* tp_flags */
722 cycle_doc, /* tp_doc */
723 (traverseproc)cycle_traverse, /* tp_traverse */
724 0, /* tp_clear */
725 0, /* tp_richcompare */
726 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000727 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728 (iternextfunc)cycle_next, /* tp_iternext */
729 0, /* tp_methods */
730 0, /* tp_members */
731 0, /* tp_getset */
732 0, /* tp_base */
733 0, /* tp_dict */
734 0, /* tp_descr_get */
735 0, /* tp_descr_set */
736 0, /* tp_dictoffset */
737 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000738 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000739 cycle_new, /* tp_new */
740 PyObject_GC_Del, /* tp_free */
741};
742
743
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000744/* dropwhile object **********************************************************/
745
746typedef struct {
747 PyObject_HEAD
748 PyObject *func;
749 PyObject *it;
750 long start;
751} dropwhileobject;
752
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000753static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000754
755static PyObject *
756dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
757{
758 PyObject *func, *seq;
759 PyObject *it;
760 dropwhileobject *lz;
761
762 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
763 return NULL;
764
765 /* Get iterator. */
766 it = PyObject_GetIter(seq);
767 if (it == NULL)
768 return NULL;
769
770 /* create dropwhileobject structure */
771 lz = (dropwhileobject *)type->tp_alloc(type, 0);
772 if (lz == NULL) {
773 Py_DECREF(it);
774 return NULL;
775 }
776 Py_INCREF(func);
777 lz->func = func;
778 lz->it = it;
779 lz->start = 0;
780
781 return (PyObject *)lz;
782}
783
784static void
785dropwhile_dealloc(dropwhileobject *lz)
786{
787 PyObject_GC_UnTrack(lz);
788 Py_XDECREF(lz->func);
789 Py_XDECREF(lz->it);
790 lz->ob_type->tp_free(lz);
791}
792
793static int
794dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
795{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000796 Py_VISIT(lz->it);
797 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000798 return 0;
799}
800
801static PyObject *
802dropwhile_next(dropwhileobject *lz)
803{
804 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000805 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000806 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000807 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000808
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000809 assert(PyIter_Check(it));
810 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000811 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000812 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000813 if (item == NULL)
814 return NULL;
815 if (lz->start == 1)
816 return item;
817
818 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
819 if (good == NULL) {
820 Py_DECREF(item);
821 return NULL;
822 }
823 ok = PyObject_IsTrue(good);
824 Py_DECREF(good);
825 if (!ok) {
826 lz->start = 1;
827 return item;
828 }
829 Py_DECREF(item);
830 }
831}
832
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000833PyDoc_STRVAR(dropwhile_doc,
834"dropwhile(predicate, iterable) --> dropwhile object\n\
835\n\
836Drop items from the iterable while predicate(item) is true.\n\
837Afterwards, return every element until the iterable is exhausted.");
838
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000839static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000840 PyObject_HEAD_INIT(NULL)
841 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000842 "itertools.dropwhile", /* tp_name */
843 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000844 0, /* tp_itemsize */
845 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000846 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000847 0, /* tp_print */
848 0, /* tp_getattr */
849 0, /* tp_setattr */
850 0, /* tp_compare */
851 0, /* tp_repr */
852 0, /* tp_as_number */
853 0, /* tp_as_sequence */
854 0, /* tp_as_mapping */
855 0, /* tp_hash */
856 0, /* tp_call */
857 0, /* tp_str */
858 PyObject_GenericGetAttr, /* tp_getattro */
859 0, /* tp_setattro */
860 0, /* tp_as_buffer */
861 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
862 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000863 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864 (traverseproc)dropwhile_traverse, /* tp_traverse */
865 0, /* tp_clear */
866 0, /* tp_richcompare */
867 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000868 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000869 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000870 0, /* tp_methods */
871 0, /* tp_members */
872 0, /* tp_getset */
873 0, /* tp_base */
874 0, /* tp_dict */
875 0, /* tp_descr_get */
876 0, /* tp_descr_set */
877 0, /* tp_dictoffset */
878 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000879 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000880 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881 PyObject_GC_Del, /* tp_free */
882};
883
884
885/* takewhile object **********************************************************/
886
887typedef struct {
888 PyObject_HEAD
889 PyObject *func;
890 PyObject *it;
891 long stop;
892} takewhileobject;
893
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000894static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000895
896static PyObject *
897takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
898{
899 PyObject *func, *seq;
900 PyObject *it;
901 takewhileobject *lz;
902
903 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
904 return NULL;
905
906 /* Get iterator. */
907 it = PyObject_GetIter(seq);
908 if (it == NULL)
909 return NULL;
910
911 /* create takewhileobject structure */
912 lz = (takewhileobject *)type->tp_alloc(type, 0);
913 if (lz == NULL) {
914 Py_DECREF(it);
915 return NULL;
916 }
917 Py_INCREF(func);
918 lz->func = func;
919 lz->it = it;
920 lz->stop = 0;
921
922 return (PyObject *)lz;
923}
924
925static void
926takewhile_dealloc(takewhileobject *lz)
927{
928 PyObject_GC_UnTrack(lz);
929 Py_XDECREF(lz->func);
930 Py_XDECREF(lz->it);
931 lz->ob_type->tp_free(lz);
932}
933
934static int
935takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
936{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000937 Py_VISIT(lz->it);
938 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000939 return 0;
940}
941
942static PyObject *
943takewhile_next(takewhileobject *lz)
944{
945 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000946 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000947 long ok;
948
949 if (lz->stop == 1)
950 return NULL;
951
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000952 assert(PyIter_Check(it));
953 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000954 if (item == NULL)
955 return NULL;
956
957 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
958 if (good == NULL) {
959 Py_DECREF(item);
960 return NULL;
961 }
962 ok = PyObject_IsTrue(good);
963 Py_DECREF(good);
964 if (ok)
965 return item;
966 Py_DECREF(item);
967 lz->stop = 1;
968 return NULL;
969}
970
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000971PyDoc_STRVAR(takewhile_doc,
972"takewhile(predicate, iterable) --> takewhile object\n\
973\n\
974Return successive entries from an iterable as long as the \n\
975predicate evaluates to true for each entry.");
976
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000977static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000978 PyObject_HEAD_INIT(NULL)
979 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000980 "itertools.takewhile", /* tp_name */
981 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000982 0, /* tp_itemsize */
983 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000984 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000985 0, /* tp_print */
986 0, /* tp_getattr */
987 0, /* tp_setattr */
988 0, /* tp_compare */
989 0, /* tp_repr */
990 0, /* tp_as_number */
991 0, /* tp_as_sequence */
992 0, /* tp_as_mapping */
993 0, /* tp_hash */
994 0, /* tp_call */
995 0, /* tp_str */
996 PyObject_GenericGetAttr, /* tp_getattro */
997 0, /* tp_setattro */
998 0, /* tp_as_buffer */
999 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1000 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001001 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001002 (traverseproc)takewhile_traverse, /* tp_traverse */
1003 0, /* tp_clear */
1004 0, /* tp_richcompare */
1005 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001006 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001007 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001008 0, /* tp_methods */
1009 0, /* tp_members */
1010 0, /* tp_getset */
1011 0, /* tp_base */
1012 0, /* tp_dict */
1013 0, /* tp_descr_get */
1014 0, /* tp_descr_set */
1015 0, /* tp_dictoffset */
1016 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001017 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001018 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001019 PyObject_GC_Del, /* tp_free */
1020};
1021
1022
1023/* islice object ************************************************************/
1024
1025typedef struct {
1026 PyObject_HEAD
1027 PyObject *it;
1028 long next;
1029 long stop;
1030 long step;
1031 long cnt;
1032} isliceobject;
1033
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001034static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001035
1036static PyObject *
1037islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1038{
1039 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001040 long start=0, stop=-1, step=1;
1041 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001042 int numargs;
1043 isliceobject *lz;
1044
1045 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +00001046 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001047 return NULL;
1048
1049 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001050 if (a1 != Py_None) {
1051 stop = PyInt_AsLong(a1);
1052 if (stop == -1) {
1053 if (PyErr_Occurred())
1054 PyErr_Clear();
1055 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001056 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001057 return NULL;
1058 }
1059 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001060 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001061 start = PyInt_AsLong(a1);
1062 if (start == -1 && PyErr_Occurred()) {
1063 PyErr_Clear();
1064 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001065 "Start argument must be a non-negative integer.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001066 return NULL;
1067 }
1068 if (a2 != Py_None) {
1069 stop = PyInt_AsLong(a2);
1070 if (stop == -1) {
1071 if (PyErr_Occurred())
1072 PyErr_Clear();
1073 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001074 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001075 return NULL;
1076 }
1077 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001078 }
1079
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001080 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001082 "Indices for islice() must be non-negative integers.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083 return NULL;
1084 }
1085
1086 if (step<1) {
1087 PyErr_SetString(PyExc_ValueError,
1088 "Step must be one or larger for islice().");
1089 return NULL;
1090 }
1091
1092 /* Get iterator. */
1093 it = PyObject_GetIter(seq);
1094 if (it == NULL)
1095 return NULL;
1096
1097 /* create isliceobject structure */
1098 lz = (isliceobject *)type->tp_alloc(type, 0);
1099 if (lz == NULL) {
1100 Py_DECREF(it);
1101 return NULL;
1102 }
1103 lz->it = it;
1104 lz->next = start;
1105 lz->stop = stop;
1106 lz->step = step;
1107 lz->cnt = 0L;
1108
1109 return (PyObject *)lz;
1110}
1111
1112static void
1113islice_dealloc(isliceobject *lz)
1114{
1115 PyObject_GC_UnTrack(lz);
1116 Py_XDECREF(lz->it);
1117 lz->ob_type->tp_free(lz);
1118}
1119
1120static int
1121islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1122{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001123 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001124 return 0;
1125}
1126
1127static PyObject *
1128islice_next(isliceobject *lz)
1129{
1130 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001131 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001132 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001133 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001134
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001135 assert(PyIter_Check(it));
1136 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001137 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001138 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001139 if (item == NULL)
1140 return NULL;
1141 Py_DECREF(item);
1142 lz->cnt++;
1143 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001144 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001145 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001146 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001147 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001148 if (item == NULL)
1149 return NULL;
1150 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001151 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001152 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001153 if (lz->next < oldnext) /* Check for overflow */
1154 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001155 return item;
1156}
1157
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001158PyDoc_STRVAR(islice_doc,
1159"islice(iterable, [start,] stop [, step]) --> islice object\n\
1160\n\
1161Return an iterator whose next() method returns selected values from an\n\
1162iterable. If start is specified, will skip all preceding elements;\n\
1163otherwise, start defaults to zero. Step defaults to one. If\n\
1164specified as another value, step determines how many values are \n\
1165skipped between successive calls. Works like a slice() on a list\n\
1166but returns an iterator.");
1167
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001168static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001169 PyObject_HEAD_INIT(NULL)
1170 0, /* ob_size */
1171 "itertools.islice", /* tp_name */
1172 sizeof(isliceobject), /* tp_basicsize */
1173 0, /* tp_itemsize */
1174 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001175 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001176 0, /* tp_print */
1177 0, /* tp_getattr */
1178 0, /* tp_setattr */
1179 0, /* tp_compare */
1180 0, /* tp_repr */
1181 0, /* tp_as_number */
1182 0, /* tp_as_sequence */
1183 0, /* tp_as_mapping */
1184 0, /* tp_hash */
1185 0, /* tp_call */
1186 0, /* tp_str */
1187 PyObject_GenericGetAttr, /* tp_getattro */
1188 0, /* tp_setattro */
1189 0, /* tp_as_buffer */
1190 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1191 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001192 islice_doc, /* tp_doc */
1193 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194 0, /* tp_clear */
1195 0, /* tp_richcompare */
1196 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001197 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001198 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001199 0, /* tp_methods */
1200 0, /* tp_members */
1201 0, /* tp_getset */
1202 0, /* tp_base */
1203 0, /* tp_dict */
1204 0, /* tp_descr_get */
1205 0, /* tp_descr_set */
1206 0, /* tp_dictoffset */
1207 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001208 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001209 islice_new, /* tp_new */
1210 PyObject_GC_Del, /* tp_free */
1211};
1212
1213
1214/* starmap object ************************************************************/
1215
1216typedef struct {
1217 PyObject_HEAD
1218 PyObject *func;
1219 PyObject *it;
1220} starmapobject;
1221
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001222static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001223
1224static PyObject *
1225starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1226{
1227 PyObject *func, *seq;
1228 PyObject *it;
1229 starmapobject *lz;
1230
1231 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1232 return NULL;
1233
1234 /* Get iterator. */
1235 it = PyObject_GetIter(seq);
1236 if (it == NULL)
1237 return NULL;
1238
1239 /* create starmapobject structure */
1240 lz = (starmapobject *)type->tp_alloc(type, 0);
1241 if (lz == NULL) {
1242 Py_DECREF(it);
1243 return NULL;
1244 }
1245 Py_INCREF(func);
1246 lz->func = func;
1247 lz->it = it;
1248
1249 return (PyObject *)lz;
1250}
1251
1252static void
1253starmap_dealloc(starmapobject *lz)
1254{
1255 PyObject_GC_UnTrack(lz);
1256 Py_XDECREF(lz->func);
1257 Py_XDECREF(lz->it);
1258 lz->ob_type->tp_free(lz);
1259}
1260
1261static int
1262starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1263{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001264 Py_VISIT(lz->it);
1265 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001266 return 0;
1267}
1268
1269static PyObject *
1270starmap_next(starmapobject *lz)
1271{
1272 PyObject *args;
1273 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001274 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001276 assert(PyIter_Check(it));
1277 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001278 if (args == NULL)
1279 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001280 if (!PyTuple_CheckExact(args)) {
1281 Py_DECREF(args);
1282 PyErr_SetString(PyExc_TypeError,
1283 "iterator must return a tuple");
1284 return NULL;
1285 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286 result = PyObject_Call(lz->func, args, NULL);
1287 Py_DECREF(args);
1288 return result;
1289}
1290
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001291PyDoc_STRVAR(starmap_doc,
1292"starmap(function, sequence) --> starmap object\n\
1293\n\
1294Return an iterator whose values are returned from the function evaluated\n\
1295with a argument tuple taken from the given sequence.");
1296
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001297static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001298 PyObject_HEAD_INIT(NULL)
1299 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001300 "itertools.starmap", /* tp_name */
1301 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001302 0, /* tp_itemsize */
1303 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001304 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001305 0, /* tp_print */
1306 0, /* tp_getattr */
1307 0, /* tp_setattr */
1308 0, /* tp_compare */
1309 0, /* tp_repr */
1310 0, /* tp_as_number */
1311 0, /* tp_as_sequence */
1312 0, /* tp_as_mapping */
1313 0, /* tp_hash */
1314 0, /* tp_call */
1315 0, /* tp_str */
1316 PyObject_GenericGetAttr, /* tp_getattro */
1317 0, /* tp_setattro */
1318 0, /* tp_as_buffer */
1319 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1320 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001321 starmap_doc, /* tp_doc */
1322 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001323 0, /* tp_clear */
1324 0, /* tp_richcompare */
1325 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001326 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001327 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328 0, /* tp_methods */
1329 0, /* tp_members */
1330 0, /* tp_getset */
1331 0, /* tp_base */
1332 0, /* tp_dict */
1333 0, /* tp_descr_get */
1334 0, /* tp_descr_set */
1335 0, /* tp_dictoffset */
1336 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001337 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001338 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339 PyObject_GC_Del, /* tp_free */
1340};
1341
1342
1343/* imap object ************************************************************/
1344
1345typedef struct {
1346 PyObject_HEAD
1347 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001348 PyObject *func;
1349} imapobject;
1350
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001351static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352
1353static PyObject *
1354imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1355{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001356 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001357 imapobject *lz;
1358 int numargs, i;
1359
1360 numargs = PyTuple_Size(args);
1361 if (numargs < 2) {
1362 PyErr_SetString(PyExc_TypeError,
1363 "imap() must have at least two arguments.");
1364 return NULL;
1365 }
1366
1367 iters = PyTuple_New(numargs-1);
1368 if (iters == NULL)
1369 return NULL;
1370
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001371 for (i=1 ; i<numargs ; i++) {
1372 /* Get iterator. */
1373 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1374 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001375 Py_DECREF(iters);
1376 return NULL;
1377 }
1378 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001379 }
1380
1381 /* create imapobject structure */
1382 lz = (imapobject *)type->tp_alloc(type, 0);
1383 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384 Py_DECREF(iters);
1385 return NULL;
1386 }
1387 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388 func = PyTuple_GET_ITEM(args, 0);
1389 Py_INCREF(func);
1390 lz->func = func;
1391
1392 return (PyObject *)lz;
1393}
1394
1395static void
1396imap_dealloc(imapobject *lz)
1397{
1398 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001399 Py_XDECREF(lz->iters);
1400 Py_XDECREF(lz->func);
1401 lz->ob_type->tp_free(lz);
1402}
1403
1404static int
1405imap_traverse(imapobject *lz, visitproc visit, void *arg)
1406{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001407 Py_VISIT(lz->iters);
1408 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001409 return 0;
1410}
1411
Raymond Hettinger2012f172003-02-07 05:32:58 +00001412/*
1413imap() is an iterator version of __builtins__.map() except that it does
1414not have the None fill-in feature. That was intentionally left out for
1415the following reasons:
1416
1417 1) Itertools are designed to be easily combined and chained together.
1418 Having all tools stop with the shortest input is a unifying principle
1419 that makes it easier to combine finite iterators (supplying data) with
1420 infinite iterators like count() and repeat() (for supplying sequential
1421 or constant arguments to a function).
1422
1423 2) In typical use cases for combining itertools, having one finite data
1424 supplier run out before another is likely to be an error condition which
1425 should not pass silently by automatically supplying None.
1426
1427 3) The use cases for automatic None fill-in are rare -- not many functions
1428 do something useful when a parameter suddenly switches type and becomes
1429 None.
1430
1431 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001432 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001433
1434 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1435*/
1436
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437static PyObject *
1438imap_next(imapobject *lz)
1439{
1440 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001441 PyObject *argtuple;
1442 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001443 int numargs, i;
1444
1445 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001446 argtuple = PyTuple_New(numargs);
1447 if (argtuple == NULL)
1448 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001449
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001450 for (i=0 ; i<numargs ; i++) {
1451 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1452 if (val == NULL) {
1453 Py_DECREF(argtuple);
1454 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001456 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001458 if (lz->func == Py_None)
1459 return argtuple;
1460 result = PyObject_Call(lz->func, argtuple, NULL);
1461 Py_DECREF(argtuple);
1462 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463}
1464
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465PyDoc_STRVAR(imap_doc,
1466"imap(func, *iterables) --> imap object\n\
1467\n\
1468Make an iterator that computes the function using arguments from\n\
1469each of the iterables. Like map() except that it returns\n\
1470an iterator instead of a list and that it stops when the shortest\n\
1471iterable is exhausted instead of filling in None for shorter\n\
1472iterables.");
1473
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001474static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475 PyObject_HEAD_INIT(NULL)
1476 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001477 "itertools.imap", /* tp_name */
1478 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479 0, /* tp_itemsize */
1480 /* methods */
1481 (destructor)imap_dealloc, /* tp_dealloc */
1482 0, /* tp_print */
1483 0, /* tp_getattr */
1484 0, /* tp_setattr */
1485 0, /* tp_compare */
1486 0, /* tp_repr */
1487 0, /* tp_as_number */
1488 0, /* tp_as_sequence */
1489 0, /* tp_as_mapping */
1490 0, /* tp_hash */
1491 0, /* tp_call */
1492 0, /* tp_str */
1493 PyObject_GenericGetAttr, /* tp_getattro */
1494 0, /* tp_setattro */
1495 0, /* tp_as_buffer */
1496 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1497 Py_TPFLAGS_BASETYPE, /* tp_flags */
1498 imap_doc, /* tp_doc */
1499 (traverseproc)imap_traverse, /* tp_traverse */
1500 0, /* tp_clear */
1501 0, /* tp_richcompare */
1502 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001503 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001504 (iternextfunc)imap_next, /* tp_iternext */
1505 0, /* tp_methods */
1506 0, /* tp_members */
1507 0, /* tp_getset */
1508 0, /* tp_base */
1509 0, /* tp_dict */
1510 0, /* tp_descr_get */
1511 0, /* tp_descr_set */
1512 0, /* tp_dictoffset */
1513 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001514 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001515 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516 PyObject_GC_Del, /* tp_free */
1517};
1518
1519
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001520/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001521
1522typedef struct {
1523 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001524 long tuplesize;
1525 long iternum; /* which iterator is active */
1526 PyObject *ittuple; /* tuple of iterators */
1527} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001529static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530
1531static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001532chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001534 chainobject *lz;
1535 int tuplesize = PySequence_Length(args);
1536 int i;
1537 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001539 /* obtain iterators */
1540 assert(PyTuple_Check(args));
1541 ittuple = PyTuple_New(tuplesize);
1542 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001544 for (i=0; i < tuplesize; ++i) {
1545 PyObject *item = PyTuple_GET_ITEM(args, i);
1546 PyObject *it = PyObject_GetIter(item);
1547 if (it == NULL) {
1548 if (PyErr_ExceptionMatches(PyExc_TypeError))
1549 PyErr_Format(PyExc_TypeError,
1550 "chain argument #%d must support iteration",
1551 i+1);
1552 Py_DECREF(ittuple);
1553 return NULL;
1554 }
1555 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556 }
1557
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001558 /* create chainobject structure */
1559 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001560 if (lz == NULL) {
1561 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001563 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001564
1565 lz->ittuple = ittuple;
1566 lz->iternum = 0;
1567 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001568
1569 return (PyObject *)lz;
1570}
1571
1572static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001573chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574{
1575 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001576 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001577 lz->ob_type->tp_free(lz);
1578}
1579
Raymond Hettinger2012f172003-02-07 05:32:58 +00001580static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001581chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001582{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001583 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001584 return 0;
1585}
1586
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001587static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001588chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001590 PyObject *it;
1591 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001592
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001593 while (lz->iternum < lz->tuplesize) {
1594 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1595 item = PyIter_Next(it);
1596 if (item != NULL)
1597 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001598 if (PyErr_Occurred()) {
1599 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1600 PyErr_Clear();
1601 else
1602 return NULL;
1603 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001604 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001605 }
1606 return NULL;
1607}
1608
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001609PyDoc_STRVAR(chain_doc,
1610"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001612Return a chain object whose .next() method returns elements from the\n\
1613first iterable until it is exhausted, then elements from the next\n\
1614iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001616static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617 PyObject_HEAD_INIT(NULL)
1618 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001619 "itertools.chain", /* tp_name */
1620 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621 0, /* tp_itemsize */
1622 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001623 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001624 0, /* tp_print */
1625 0, /* tp_getattr */
1626 0, /* tp_setattr */
1627 0, /* tp_compare */
1628 0, /* tp_repr */
1629 0, /* tp_as_number */
1630 0, /* tp_as_sequence */
1631 0, /* tp_as_mapping */
1632 0, /* tp_hash */
1633 0, /* tp_call */
1634 0, /* tp_str */
1635 PyObject_GenericGetAttr, /* tp_getattro */
1636 0, /* tp_setattro */
1637 0, /* tp_as_buffer */
1638 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1639 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001640 chain_doc, /* tp_doc */
1641 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001642 0, /* tp_clear */
1643 0, /* tp_richcompare */
1644 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001645 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001646 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001647 0, /* tp_methods */
1648 0, /* tp_members */
1649 0, /* tp_getset */
1650 0, /* tp_base */
1651 0, /* tp_dict */
1652 0, /* tp_descr_get */
1653 0, /* tp_descr_set */
1654 0, /* tp_dictoffset */
1655 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001656 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001657 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658 PyObject_GC_Del, /* tp_free */
1659};
1660
1661
1662/* ifilter object ************************************************************/
1663
1664typedef struct {
1665 PyObject_HEAD
1666 PyObject *func;
1667 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668} ifilterobject;
1669
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001670static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671
1672static PyObject *
1673ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1674{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001675 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676 PyObject *it;
1677 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001678
Raymond Hettinger60eca932003-02-09 06:40:58 +00001679 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680 return NULL;
1681
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682 /* Get iterator. */
1683 it = PyObject_GetIter(seq);
1684 if (it == NULL)
1685 return NULL;
1686
1687 /* create ifilterobject structure */
1688 lz = (ifilterobject *)type->tp_alloc(type, 0);
1689 if (lz == NULL) {
1690 Py_DECREF(it);
1691 return NULL;
1692 }
1693 Py_INCREF(func);
1694 lz->func = func;
1695 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001696
1697 return (PyObject *)lz;
1698}
1699
1700static void
1701ifilter_dealloc(ifilterobject *lz)
1702{
1703 PyObject_GC_UnTrack(lz);
1704 Py_XDECREF(lz->func);
1705 Py_XDECREF(lz->it);
1706 lz->ob_type->tp_free(lz);
1707}
1708
1709static int
1710ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1711{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001712 Py_VISIT(lz->it);
1713 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001714 return 0;
1715}
1716
1717static PyObject *
1718ifilter_next(ifilterobject *lz)
1719{
1720 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001721 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001722 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001723 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001725 assert(PyIter_Check(it));
1726 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001728 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001729 if (item == NULL)
1730 return NULL;
1731
1732 if (lz->func == Py_None) {
1733 ok = PyObject_IsTrue(item);
1734 } else {
1735 PyObject *good;
1736 good = PyObject_CallFunctionObjArgs(lz->func,
1737 item, NULL);
1738 if (good == NULL) {
1739 Py_DECREF(item);
1740 return NULL;
1741 }
1742 ok = PyObject_IsTrue(good);
1743 Py_DECREF(good);
1744 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001745 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746 return item;
1747 Py_DECREF(item);
1748 }
1749}
1750
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001751PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001752"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001754Return those items of sequence for which function(item) is true.\n\
1755If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001756
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001757static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758 PyObject_HEAD_INIT(NULL)
1759 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001760 "itertools.ifilter", /* tp_name */
1761 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001762 0, /* tp_itemsize */
1763 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001764 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001765 0, /* tp_print */
1766 0, /* tp_getattr */
1767 0, /* tp_setattr */
1768 0, /* tp_compare */
1769 0, /* tp_repr */
1770 0, /* tp_as_number */
1771 0, /* tp_as_sequence */
1772 0, /* tp_as_mapping */
1773 0, /* tp_hash */
1774 0, /* tp_call */
1775 0, /* tp_str */
1776 PyObject_GenericGetAttr, /* tp_getattro */
1777 0, /* tp_setattro */
1778 0, /* tp_as_buffer */
1779 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1780 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001781 ifilter_doc, /* tp_doc */
1782 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001783 0, /* tp_clear */
1784 0, /* tp_richcompare */
1785 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001786 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001787 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001788 0, /* tp_methods */
1789 0, /* tp_members */
1790 0, /* tp_getset */
1791 0, /* tp_base */
1792 0, /* tp_dict */
1793 0, /* tp_descr_get */
1794 0, /* tp_descr_set */
1795 0, /* tp_dictoffset */
1796 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001797 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001798 ifilter_new, /* tp_new */
1799 PyObject_GC_Del, /* tp_free */
1800};
1801
1802
1803/* ifilterfalse object ************************************************************/
1804
1805typedef struct {
1806 PyObject_HEAD
1807 PyObject *func;
1808 PyObject *it;
1809} ifilterfalseobject;
1810
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001811static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001812
1813static PyObject *
1814ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1815{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001816 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001817 PyObject *it;
1818 ifilterfalseobject *lz;
1819
1820 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1821 return NULL;
1822
1823 /* Get iterator. */
1824 it = PyObject_GetIter(seq);
1825 if (it == NULL)
1826 return NULL;
1827
1828 /* create ifilterfalseobject structure */
1829 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1830 if (lz == NULL) {
1831 Py_DECREF(it);
1832 return NULL;
1833 }
1834 Py_INCREF(func);
1835 lz->func = func;
1836 lz->it = it;
1837
1838 return (PyObject *)lz;
1839}
1840
1841static void
1842ifilterfalse_dealloc(ifilterfalseobject *lz)
1843{
1844 PyObject_GC_UnTrack(lz);
1845 Py_XDECREF(lz->func);
1846 Py_XDECREF(lz->it);
1847 lz->ob_type->tp_free(lz);
1848}
1849
1850static int
1851ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1852{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001853 Py_VISIT(lz->it);
1854 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001855 return 0;
1856}
1857
1858static PyObject *
1859ifilterfalse_next(ifilterfalseobject *lz)
1860{
1861 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001862 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001863 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001864 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001865
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001866 assert(PyIter_Check(it));
1867 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001868 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001869 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001870 if (item == NULL)
1871 return NULL;
1872
1873 if (lz->func == Py_None) {
1874 ok = PyObject_IsTrue(item);
1875 } else {
1876 PyObject *good;
1877 good = PyObject_CallFunctionObjArgs(lz->func,
1878 item, NULL);
1879 if (good == NULL) {
1880 Py_DECREF(item);
1881 return NULL;
1882 }
1883 ok = PyObject_IsTrue(good);
1884 Py_DECREF(good);
1885 }
1886 if (!ok)
1887 return item;
1888 Py_DECREF(item);
1889 }
1890}
1891
Raymond Hettinger60eca932003-02-09 06:40:58 +00001892PyDoc_STRVAR(ifilterfalse_doc,
1893"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1894\n\
1895Return those items of sequence for which function(item) is false.\n\
1896If function is None, return the items that are false.");
1897
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001898static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001899 PyObject_HEAD_INIT(NULL)
1900 0, /* ob_size */
1901 "itertools.ifilterfalse", /* tp_name */
1902 sizeof(ifilterfalseobject), /* tp_basicsize */
1903 0, /* tp_itemsize */
1904 /* methods */
1905 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1906 0, /* tp_print */
1907 0, /* tp_getattr */
1908 0, /* tp_setattr */
1909 0, /* tp_compare */
1910 0, /* tp_repr */
1911 0, /* tp_as_number */
1912 0, /* tp_as_sequence */
1913 0, /* tp_as_mapping */
1914 0, /* tp_hash */
1915 0, /* tp_call */
1916 0, /* tp_str */
1917 PyObject_GenericGetAttr, /* tp_getattro */
1918 0, /* tp_setattro */
1919 0, /* tp_as_buffer */
1920 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1921 Py_TPFLAGS_BASETYPE, /* tp_flags */
1922 ifilterfalse_doc, /* tp_doc */
1923 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1924 0, /* tp_clear */
1925 0, /* tp_richcompare */
1926 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001927 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001928 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1929 0, /* tp_methods */
1930 0, /* tp_members */
1931 0, /* tp_getset */
1932 0, /* tp_base */
1933 0, /* tp_dict */
1934 0, /* tp_descr_get */
1935 0, /* tp_descr_set */
1936 0, /* tp_dictoffset */
1937 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001938 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001939 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001940 PyObject_GC_Del, /* tp_free */
1941};
1942
1943
1944/* count object ************************************************************/
1945
1946typedef struct {
1947 PyObject_HEAD
1948 long cnt;
1949} countobject;
1950
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001951static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001952
1953static PyObject *
1954count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1955{
1956 countobject *lz;
1957 long cnt = 0;
1958
1959 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1960 return NULL;
1961
1962 /* create countobject structure */
1963 lz = (countobject *)PyObject_New(countobject, &count_type);
1964 if (lz == NULL)
1965 return NULL;
1966 lz->cnt = cnt;
1967
1968 return (PyObject *)lz;
1969}
1970
1971static PyObject *
1972count_next(countobject *lz)
1973{
1974 return PyInt_FromLong(lz->cnt++);
1975}
1976
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001977static PyObject *
1978count_repr(countobject *lz)
1979{
Brett Cannon00468392004-04-13 02:43:53 +00001980 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001981}
1982
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001983PyDoc_STRVAR(count_doc,
1984"count([firstval]) --> count object\n\
1985\n\
1986Return a count object whose .next() method returns consecutive\n\
1987integers starting from zero or, if specified, from firstval.");
1988
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001989static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001990 PyObject_HEAD_INIT(NULL)
1991 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001992 "itertools.count", /* tp_name */
1993 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001994 0, /* tp_itemsize */
1995 /* methods */
1996 (destructor)PyObject_Del, /* tp_dealloc */
1997 0, /* tp_print */
1998 0, /* tp_getattr */
1999 0, /* tp_setattr */
2000 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002001 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002002 0, /* tp_as_number */
2003 0, /* tp_as_sequence */
2004 0, /* tp_as_mapping */
2005 0, /* tp_hash */
2006 0, /* tp_call */
2007 0, /* tp_str */
2008 PyObject_GenericGetAttr, /* tp_getattro */
2009 0, /* tp_setattro */
2010 0, /* tp_as_buffer */
2011 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002012 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002013 0, /* tp_traverse */
2014 0, /* tp_clear */
2015 0, /* tp_richcompare */
2016 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002017 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002018 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002019 0, /* tp_methods */
2020 0, /* tp_members */
2021 0, /* tp_getset */
2022 0, /* tp_base */
2023 0, /* tp_dict */
2024 0, /* tp_descr_get */
2025 0, /* tp_descr_set */
2026 0, /* tp_dictoffset */
2027 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002028 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002029 count_new, /* tp_new */
2030};
2031
2032
2033/* izip object ************************************************************/
2034
2035#include "Python.h"
2036
2037typedef struct {
2038 PyObject_HEAD
2039 long tuplesize;
2040 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002041 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002042} izipobject;
2043
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002044static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002045
2046static PyObject *
2047izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2048{
2049 izipobject *lz;
2050 int i;
2051 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002052 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002053 int tuplesize = PySequence_Length(args);
2054
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002055 /* args must be a tuple */
2056 assert(PyTuple_Check(args));
2057
2058 /* obtain iterators */
2059 ittuple = PyTuple_New(tuplesize);
2060 if(ittuple == NULL)
2061 return NULL;
2062 for (i=0; i < tuplesize; ++i) {
2063 PyObject *item = PyTuple_GET_ITEM(args, i);
2064 PyObject *it = PyObject_GetIter(item);
2065 if (it == NULL) {
2066 if (PyErr_ExceptionMatches(PyExc_TypeError))
2067 PyErr_Format(PyExc_TypeError,
2068 "izip argument #%d must support iteration",
2069 i+1);
2070 Py_DECREF(ittuple);
2071 return NULL;
2072 }
2073 PyTuple_SET_ITEM(ittuple, i, it);
2074 }
2075
Raymond Hettinger2012f172003-02-07 05:32:58 +00002076 /* create a result holder */
2077 result = PyTuple_New(tuplesize);
2078 if (result == NULL) {
2079 Py_DECREF(ittuple);
2080 return NULL;
2081 }
2082 for (i=0 ; i < tuplesize ; i++) {
2083 Py_INCREF(Py_None);
2084 PyTuple_SET_ITEM(result, i, Py_None);
2085 }
2086
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002087 /* create izipobject structure */
2088 lz = (izipobject *)type->tp_alloc(type, 0);
2089 if (lz == NULL) {
2090 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002091 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002092 return NULL;
2093 }
2094 lz->ittuple = ittuple;
2095 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002096 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002097
2098 return (PyObject *)lz;
2099}
2100
2101static void
2102izip_dealloc(izipobject *lz)
2103{
2104 PyObject_GC_UnTrack(lz);
2105 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002106 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002107 lz->ob_type->tp_free(lz);
2108}
2109
2110static int
2111izip_traverse(izipobject *lz, visitproc visit, void *arg)
2112{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002113 Py_VISIT(lz->ittuple);
2114 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002115 return 0;
2116}
2117
2118static PyObject *
2119izip_next(izipobject *lz)
2120{
2121 int i;
2122 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002123 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002124 PyObject *it;
2125 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002126 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002127
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002128 if (tuplesize == 0)
2129 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002130 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002131 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002132 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002133 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002134 assert(PyIter_Check(it));
2135 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002136 if (item == NULL) {
2137 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002138 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002139 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002140 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002141 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002142 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002143 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002144 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002145 result = PyTuple_New(tuplesize);
2146 if (result == NULL)
2147 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002148 for (i=0 ; i < tuplesize ; i++) {
2149 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002150 assert(PyIter_Check(it));
2151 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002152 if (item == NULL) {
2153 Py_DECREF(result);
2154 return NULL;
2155 }
2156 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002157 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002158 }
2159 return result;
2160}
2161
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002162PyDoc_STRVAR(izip_doc,
2163"izip(iter1 [,iter2 [...]]) --> izip object\n\
2164\n\
2165Return a izip object whose .next() method returns a tuple where\n\
2166the i-th element comes from the i-th iterable argument. The .next()\n\
2167method continues until the shortest iterable in the argument sequence\n\
2168is exhausted and then it raises StopIteration. Works like the zip()\n\
2169function but consumes less memory by returning an iterator instead of\n\
2170a list.");
2171
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002172static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002173 PyObject_HEAD_INIT(NULL)
2174 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002175 "itertools.izip", /* tp_name */
2176 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002177 0, /* tp_itemsize */
2178 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002179 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002180 0, /* tp_print */
2181 0, /* tp_getattr */
2182 0, /* tp_setattr */
2183 0, /* tp_compare */
2184 0, /* tp_repr */
2185 0, /* tp_as_number */
2186 0, /* tp_as_sequence */
2187 0, /* tp_as_mapping */
2188 0, /* tp_hash */
2189 0, /* tp_call */
2190 0, /* tp_str */
2191 PyObject_GenericGetAttr, /* tp_getattro */
2192 0, /* tp_setattro */
2193 0, /* tp_as_buffer */
2194 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2195 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002196 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002197 (traverseproc)izip_traverse, /* tp_traverse */
2198 0, /* tp_clear */
2199 0, /* tp_richcompare */
2200 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002201 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002202 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002203 0, /* tp_methods */
2204 0, /* tp_members */
2205 0, /* tp_getset */
2206 0, /* tp_base */
2207 0, /* tp_dict */
2208 0, /* tp_descr_get */
2209 0, /* tp_descr_set */
2210 0, /* tp_dictoffset */
2211 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002212 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002213 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002214 PyObject_GC_Del, /* tp_free */
2215};
2216
2217
2218/* repeat object ************************************************************/
2219
2220typedef struct {
2221 PyObject_HEAD
2222 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002223 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002224} repeatobject;
2225
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002226static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002227
2228static PyObject *
2229repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2230{
2231 repeatobject *ro;
2232 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002233 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002234
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002235 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002236 return NULL;
2237
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002238 if (PyTuple_Size(args) == 2 && cnt < 0)
2239 cnt = 0;
2240
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002241 ro = (repeatobject *)type->tp_alloc(type, 0);
2242 if (ro == NULL)
2243 return NULL;
2244 Py_INCREF(element);
2245 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002246 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002247 return (PyObject *)ro;
2248}
2249
2250static void
2251repeat_dealloc(repeatobject *ro)
2252{
2253 PyObject_GC_UnTrack(ro);
2254 Py_XDECREF(ro->element);
2255 ro->ob_type->tp_free(ro);
2256}
2257
2258static int
2259repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2260{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002261 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002262 return 0;
2263}
2264
2265static PyObject *
2266repeat_next(repeatobject *ro)
2267{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002268 if (ro->cnt == 0)
2269 return NULL;
2270 if (ro->cnt > 0)
2271 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002272 Py_INCREF(ro->element);
2273 return ro->element;
2274}
2275
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002276static PyObject *
2277repeat_repr(repeatobject *ro)
2278{
2279 PyObject *result, *objrepr;
2280
2281 objrepr = PyObject_Repr(ro->element);
2282 if (objrepr == NULL)
2283 return NULL;
2284
2285 if (ro->cnt == -1)
2286 result = PyString_FromFormat("repeat(%s)",
2287 PyString_AS_STRING(objrepr));
2288 else
Brett Cannon00468392004-04-13 02:43:53 +00002289 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002290 PyString_AS_STRING(objrepr), ro->cnt);
2291 Py_DECREF(objrepr);
2292 return result;
2293}
2294
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002295static int
2296repeat_len(repeatobject *ro)
2297{
2298 if (ro->cnt == -1)
2299 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2300 return (int)(ro->cnt);
2301}
2302
2303static PySequenceMethods repeat_as_sequence = {
2304 (inquiry)repeat_len, /* sq_length */
2305 0, /* sq_concat */
2306};
2307
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002308PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002309"repeat(element [,times]) -> create an iterator which returns the element\n\
2310for the specified number of times. If not specified, returns the element\n\
2311endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002312
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002313static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002314 PyObject_HEAD_INIT(NULL)
2315 0, /* ob_size */
2316 "itertools.repeat", /* tp_name */
2317 sizeof(repeatobject), /* tp_basicsize */
2318 0, /* tp_itemsize */
2319 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002320 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002321 0, /* tp_print */
2322 0, /* tp_getattr */
2323 0, /* tp_setattr */
2324 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002325 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002326 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002327 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002328 0, /* tp_as_mapping */
2329 0, /* tp_hash */
2330 0, /* tp_call */
2331 0, /* tp_str */
2332 PyObject_GenericGetAttr, /* tp_getattro */
2333 0, /* tp_setattro */
2334 0, /* tp_as_buffer */
2335 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2336 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002337 repeat_doc, /* tp_doc */
2338 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002339 0, /* tp_clear */
2340 0, /* tp_richcompare */
2341 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002342 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002343 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002344 0, /* tp_methods */
2345 0, /* tp_members */
2346 0, /* tp_getset */
2347 0, /* tp_base */
2348 0, /* tp_dict */
2349 0, /* tp_descr_get */
2350 0, /* tp_descr_set */
2351 0, /* tp_dictoffset */
2352 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002353 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002354 repeat_new, /* tp_new */
2355 PyObject_GC_Del, /* tp_free */
2356};
2357
2358
2359/* module level code ********************************************************/
2360
2361PyDoc_STRVAR(module_doc,
2362"Functional tools for creating and using iterators.\n\
2363\n\
2364Infinite iterators:\n\
2365count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002366cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002367repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002368\n\
2369Iterators terminating on the shortest input sequence:\n\
2370izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002371ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2372ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002373islice(seq, [start,] stop [, step]) --> elements from\n\
2374 seq[start:stop:step]\n\
2375imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2376starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002377tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002378chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002379takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2380dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002381groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002382");
2383
2384
Raymond Hettingerad983e72003-11-12 14:32:26 +00002385static PyMethodDef module_methods[] = {
2386 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2387 {NULL, NULL} /* sentinel */
2388};
2389
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002390PyMODINIT_FUNC
2391inititertools(void)
2392{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002393 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002394 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002395 char *name;
2396 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002397 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002398 &dropwhile_type,
2399 &takewhile_type,
2400 &islice_type,
2401 &starmap_type,
2402 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002403 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002404 &ifilter_type,
2405 &ifilterfalse_type,
2406 &count_type,
2407 &izip_type,
2408 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002409 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002410 NULL
2411 };
2412
Skip Montanarof3938fd2004-02-10 20:27:40 +00002413 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002414 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002415
Raymond Hettinger60eca932003-02-09 06:40:58 +00002416 for (i=0 ; typelist[i] != NULL ; i++) {
2417 if (PyType_Ready(typelist[i]) < 0)
2418 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002419 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002420 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002421 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002422 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002423 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002424
2425 if (PyType_Ready(&teedataobject_type) < 0)
2426 return;
2427 if (PyType_Ready(&tee_type) < 0)
2428 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002429 if (PyType_Ready(&_grouper_type) < 0)
2430 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002431}