blob: 3da0258eefb21d4edb4207bf0efb95822feb26bf [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;
665
666 while (1) {
667 item = PyIter_Next(lz->it);
668 if (item != NULL) {
669 if (!lz->firstpass)
670 PyList_Append(lz->saved, item);
671 return item;
672 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000673 if (PyErr_Occurred()) {
674 if (PyErr_ExceptionMatches(PyExc_StopIteration))
675 PyErr_Clear();
676 else
677 return NULL;
678 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000679 if (PyList_Size(lz->saved) == 0)
680 return NULL;
681 it = PyObject_GetIter(lz->saved);
682 if (it == NULL)
683 return NULL;
684 Py_DECREF(lz->it);
685 lz->it = it;
686 lz->firstpass = 1;
687 }
688}
689
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000690PyDoc_STRVAR(cycle_doc,
691"cycle(iterable) --> cycle object\n\
692\n\
693Return elements from the iterable until it is exhausted.\n\
694Then repeat the sequence indefinitely.");
695
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000696static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000697 PyObject_HEAD_INIT(NULL)
698 0, /* ob_size */
699 "itertools.cycle", /* tp_name */
700 sizeof(cycleobject), /* tp_basicsize */
701 0, /* tp_itemsize */
702 /* methods */
703 (destructor)cycle_dealloc, /* tp_dealloc */
704 0, /* tp_print */
705 0, /* tp_getattr */
706 0, /* tp_setattr */
707 0, /* tp_compare */
708 0, /* tp_repr */
709 0, /* tp_as_number */
710 0, /* tp_as_sequence */
711 0, /* tp_as_mapping */
712 0, /* tp_hash */
713 0, /* tp_call */
714 0, /* tp_str */
715 PyObject_GenericGetAttr, /* tp_getattro */
716 0, /* tp_setattro */
717 0, /* tp_as_buffer */
718 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
719 Py_TPFLAGS_BASETYPE, /* tp_flags */
720 cycle_doc, /* tp_doc */
721 (traverseproc)cycle_traverse, /* tp_traverse */
722 0, /* tp_clear */
723 0, /* tp_richcompare */
724 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000725 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000726 (iternextfunc)cycle_next, /* tp_iternext */
727 0, /* tp_methods */
728 0, /* tp_members */
729 0, /* tp_getset */
730 0, /* tp_base */
731 0, /* tp_dict */
732 0, /* tp_descr_get */
733 0, /* tp_descr_set */
734 0, /* tp_dictoffset */
735 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000736 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000737 cycle_new, /* tp_new */
738 PyObject_GC_Del, /* tp_free */
739};
740
741
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000742/* dropwhile object **********************************************************/
743
744typedef struct {
745 PyObject_HEAD
746 PyObject *func;
747 PyObject *it;
748 long start;
749} dropwhileobject;
750
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000751static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000752
753static PyObject *
754dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
755{
756 PyObject *func, *seq;
757 PyObject *it;
758 dropwhileobject *lz;
759
760 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
761 return NULL;
762
763 /* Get iterator. */
764 it = PyObject_GetIter(seq);
765 if (it == NULL)
766 return NULL;
767
768 /* create dropwhileobject structure */
769 lz = (dropwhileobject *)type->tp_alloc(type, 0);
770 if (lz == NULL) {
771 Py_DECREF(it);
772 return NULL;
773 }
774 Py_INCREF(func);
775 lz->func = func;
776 lz->it = it;
777 lz->start = 0;
778
779 return (PyObject *)lz;
780}
781
782static void
783dropwhile_dealloc(dropwhileobject *lz)
784{
785 PyObject_GC_UnTrack(lz);
786 Py_XDECREF(lz->func);
787 Py_XDECREF(lz->it);
788 lz->ob_type->tp_free(lz);
789}
790
791static int
792dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
793{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000794 Py_VISIT(lz->it);
795 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000796 return 0;
797}
798
799static PyObject *
800dropwhile_next(dropwhileobject *lz)
801{
802 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000803 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000804 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000805 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000806
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000807 assert(PyIter_Check(it));
808 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000809 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000810 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000811 if (item == NULL)
812 return NULL;
813 if (lz->start == 1)
814 return item;
815
816 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
817 if (good == NULL) {
818 Py_DECREF(item);
819 return NULL;
820 }
821 ok = PyObject_IsTrue(good);
822 Py_DECREF(good);
823 if (!ok) {
824 lz->start = 1;
825 return item;
826 }
827 Py_DECREF(item);
828 }
829}
830
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000831PyDoc_STRVAR(dropwhile_doc,
832"dropwhile(predicate, iterable) --> dropwhile object\n\
833\n\
834Drop items from the iterable while predicate(item) is true.\n\
835Afterwards, return every element until the iterable is exhausted.");
836
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000837static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000838 PyObject_HEAD_INIT(NULL)
839 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000840 "itertools.dropwhile", /* tp_name */
841 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 0, /* tp_itemsize */
843 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000844 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000845 0, /* tp_print */
846 0, /* tp_getattr */
847 0, /* tp_setattr */
848 0, /* tp_compare */
849 0, /* tp_repr */
850 0, /* tp_as_number */
851 0, /* tp_as_sequence */
852 0, /* tp_as_mapping */
853 0, /* tp_hash */
854 0, /* tp_call */
855 0, /* tp_str */
856 PyObject_GenericGetAttr, /* tp_getattro */
857 0, /* tp_setattro */
858 0, /* tp_as_buffer */
859 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
860 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000861 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862 (traverseproc)dropwhile_traverse, /* tp_traverse */
863 0, /* tp_clear */
864 0, /* tp_richcompare */
865 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000866 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000867 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000868 0, /* tp_methods */
869 0, /* tp_members */
870 0, /* tp_getset */
871 0, /* tp_base */
872 0, /* tp_dict */
873 0, /* tp_descr_get */
874 0, /* tp_descr_set */
875 0, /* tp_dictoffset */
876 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000877 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000878 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000879 PyObject_GC_Del, /* tp_free */
880};
881
882
883/* takewhile object **********************************************************/
884
885typedef struct {
886 PyObject_HEAD
887 PyObject *func;
888 PyObject *it;
889 long stop;
890} takewhileobject;
891
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000892static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000893
894static PyObject *
895takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
896{
897 PyObject *func, *seq;
898 PyObject *it;
899 takewhileobject *lz;
900
901 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
902 return NULL;
903
904 /* Get iterator. */
905 it = PyObject_GetIter(seq);
906 if (it == NULL)
907 return NULL;
908
909 /* create takewhileobject structure */
910 lz = (takewhileobject *)type->tp_alloc(type, 0);
911 if (lz == NULL) {
912 Py_DECREF(it);
913 return NULL;
914 }
915 Py_INCREF(func);
916 lz->func = func;
917 lz->it = it;
918 lz->stop = 0;
919
920 return (PyObject *)lz;
921}
922
923static void
924takewhile_dealloc(takewhileobject *lz)
925{
926 PyObject_GC_UnTrack(lz);
927 Py_XDECREF(lz->func);
928 Py_XDECREF(lz->it);
929 lz->ob_type->tp_free(lz);
930}
931
932static int
933takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
934{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000935 Py_VISIT(lz->it);
936 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000937 return 0;
938}
939
940static PyObject *
941takewhile_next(takewhileobject *lz)
942{
943 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000944 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000945 long ok;
946
947 if (lz->stop == 1)
948 return NULL;
949
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000950 assert(PyIter_Check(it));
951 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000952 if (item == NULL)
953 return NULL;
954
955 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
956 if (good == NULL) {
957 Py_DECREF(item);
958 return NULL;
959 }
960 ok = PyObject_IsTrue(good);
961 Py_DECREF(good);
962 if (ok)
963 return item;
964 Py_DECREF(item);
965 lz->stop = 1;
966 return NULL;
967}
968
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000969PyDoc_STRVAR(takewhile_doc,
970"takewhile(predicate, iterable) --> takewhile object\n\
971\n\
972Return successive entries from an iterable as long as the \n\
973predicate evaluates to true for each entry.");
974
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000975static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000976 PyObject_HEAD_INIT(NULL)
977 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000978 "itertools.takewhile", /* tp_name */
979 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000980 0, /* tp_itemsize */
981 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000982 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000983 0, /* tp_print */
984 0, /* tp_getattr */
985 0, /* tp_setattr */
986 0, /* tp_compare */
987 0, /* tp_repr */
988 0, /* tp_as_number */
989 0, /* tp_as_sequence */
990 0, /* tp_as_mapping */
991 0, /* tp_hash */
992 0, /* tp_call */
993 0, /* tp_str */
994 PyObject_GenericGetAttr, /* tp_getattro */
995 0, /* tp_setattro */
996 0, /* tp_as_buffer */
997 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
998 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000999 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001000 (traverseproc)takewhile_traverse, /* tp_traverse */
1001 0, /* tp_clear */
1002 0, /* tp_richcompare */
1003 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001004 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001005 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006 0, /* tp_methods */
1007 0, /* tp_members */
1008 0, /* tp_getset */
1009 0, /* tp_base */
1010 0, /* tp_dict */
1011 0, /* tp_descr_get */
1012 0, /* tp_descr_set */
1013 0, /* tp_dictoffset */
1014 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001015 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001016 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001017 PyObject_GC_Del, /* tp_free */
1018};
1019
1020
1021/* islice object ************************************************************/
1022
1023typedef struct {
1024 PyObject_HEAD
1025 PyObject *it;
1026 long next;
1027 long stop;
1028 long step;
1029 long cnt;
1030} isliceobject;
1031
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001032static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001033
1034static PyObject *
1035islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1036{
1037 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001038 long start=0, stop=-1, step=1;
1039 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001040 int numargs;
1041 isliceobject *lz;
1042
1043 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +00001044 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001045 return NULL;
1046
1047 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001048 if (a1 != Py_None) {
1049 stop = PyInt_AsLong(a1);
1050 if (stop == -1) {
1051 if (PyErr_Occurred())
1052 PyErr_Clear();
1053 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001054 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001055 return NULL;
1056 }
1057 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001058 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001059 start = PyInt_AsLong(a1);
1060 if (start == -1 && PyErr_Occurred()) {
1061 PyErr_Clear();
1062 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001063 "Start argument must be a non-negative integer.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001064 return NULL;
1065 }
1066 if (a2 != Py_None) {
1067 stop = PyInt_AsLong(a2);
1068 if (stop == -1) {
1069 if (PyErr_Occurred())
1070 PyErr_Clear();
1071 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001072 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001073 return NULL;
1074 }
1075 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001076 }
1077
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001078 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001079 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001080 "Indices for islice() must be non-negative integers.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 return NULL;
1082 }
1083
1084 if (step<1) {
1085 PyErr_SetString(PyExc_ValueError,
1086 "Step must be one or larger for islice().");
1087 return NULL;
1088 }
1089
1090 /* Get iterator. */
1091 it = PyObject_GetIter(seq);
1092 if (it == NULL)
1093 return NULL;
1094
1095 /* create isliceobject structure */
1096 lz = (isliceobject *)type->tp_alloc(type, 0);
1097 if (lz == NULL) {
1098 Py_DECREF(it);
1099 return NULL;
1100 }
1101 lz->it = it;
1102 lz->next = start;
1103 lz->stop = stop;
1104 lz->step = step;
1105 lz->cnt = 0L;
1106
1107 return (PyObject *)lz;
1108}
1109
1110static void
1111islice_dealloc(isliceobject *lz)
1112{
1113 PyObject_GC_UnTrack(lz);
1114 Py_XDECREF(lz->it);
1115 lz->ob_type->tp_free(lz);
1116}
1117
1118static int
1119islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1120{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001121 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001122 return 0;
1123}
1124
1125static PyObject *
1126islice_next(isliceobject *lz)
1127{
1128 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001129 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001130 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001131 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001132
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001133 assert(PyIter_Check(it));
1134 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001136 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001137 if (item == NULL)
1138 return NULL;
1139 Py_DECREF(item);
1140 lz->cnt++;
1141 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001142 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001143 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001144 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001145 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001146 if (item == NULL)
1147 return NULL;
1148 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001149 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001150 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001151 if (lz->next < oldnext) /* Check for overflow */
1152 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153 return item;
1154}
1155
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001156PyDoc_STRVAR(islice_doc,
1157"islice(iterable, [start,] stop [, step]) --> islice object\n\
1158\n\
1159Return an iterator whose next() method returns selected values from an\n\
1160iterable. If start is specified, will skip all preceding elements;\n\
1161otherwise, start defaults to zero. Step defaults to one. If\n\
1162specified as another value, step determines how many values are \n\
1163skipped between successive calls. Works like a slice() on a list\n\
1164but returns an iterator.");
1165
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001166static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001167 PyObject_HEAD_INIT(NULL)
1168 0, /* ob_size */
1169 "itertools.islice", /* tp_name */
1170 sizeof(isliceobject), /* tp_basicsize */
1171 0, /* tp_itemsize */
1172 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001173 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001174 0, /* tp_print */
1175 0, /* tp_getattr */
1176 0, /* tp_setattr */
1177 0, /* tp_compare */
1178 0, /* tp_repr */
1179 0, /* tp_as_number */
1180 0, /* tp_as_sequence */
1181 0, /* tp_as_mapping */
1182 0, /* tp_hash */
1183 0, /* tp_call */
1184 0, /* tp_str */
1185 PyObject_GenericGetAttr, /* tp_getattro */
1186 0, /* tp_setattro */
1187 0, /* tp_as_buffer */
1188 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1189 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001190 islice_doc, /* tp_doc */
1191 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001192 0, /* tp_clear */
1193 0, /* tp_richcompare */
1194 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001195 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001196 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001197 0, /* tp_methods */
1198 0, /* tp_members */
1199 0, /* tp_getset */
1200 0, /* tp_base */
1201 0, /* tp_dict */
1202 0, /* tp_descr_get */
1203 0, /* tp_descr_set */
1204 0, /* tp_dictoffset */
1205 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001206 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207 islice_new, /* tp_new */
1208 PyObject_GC_Del, /* tp_free */
1209};
1210
1211
1212/* starmap object ************************************************************/
1213
1214typedef struct {
1215 PyObject_HEAD
1216 PyObject *func;
1217 PyObject *it;
1218} starmapobject;
1219
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001220static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
1222static PyObject *
1223starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1224{
1225 PyObject *func, *seq;
1226 PyObject *it;
1227 starmapobject *lz;
1228
1229 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1230 return NULL;
1231
1232 /* Get iterator. */
1233 it = PyObject_GetIter(seq);
1234 if (it == NULL)
1235 return NULL;
1236
1237 /* create starmapobject structure */
1238 lz = (starmapobject *)type->tp_alloc(type, 0);
1239 if (lz == NULL) {
1240 Py_DECREF(it);
1241 return NULL;
1242 }
1243 Py_INCREF(func);
1244 lz->func = func;
1245 lz->it = it;
1246
1247 return (PyObject *)lz;
1248}
1249
1250static void
1251starmap_dealloc(starmapobject *lz)
1252{
1253 PyObject_GC_UnTrack(lz);
1254 Py_XDECREF(lz->func);
1255 Py_XDECREF(lz->it);
1256 lz->ob_type->tp_free(lz);
1257}
1258
1259static int
1260starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1261{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001262 Py_VISIT(lz->it);
1263 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001264 return 0;
1265}
1266
1267static PyObject *
1268starmap_next(starmapobject *lz)
1269{
1270 PyObject *args;
1271 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001272 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001273
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001274 assert(PyIter_Check(it));
1275 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001276 if (args == NULL)
1277 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001278 if (!PyTuple_CheckExact(args)) {
1279 Py_DECREF(args);
1280 PyErr_SetString(PyExc_TypeError,
1281 "iterator must return a tuple");
1282 return NULL;
1283 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284 result = PyObject_Call(lz->func, args, NULL);
1285 Py_DECREF(args);
1286 return result;
1287}
1288
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001289PyDoc_STRVAR(starmap_doc,
1290"starmap(function, sequence) --> starmap object\n\
1291\n\
1292Return an iterator whose values are returned from the function evaluated\n\
1293with a argument tuple taken from the given sequence.");
1294
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001295static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001296 PyObject_HEAD_INIT(NULL)
1297 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001298 "itertools.starmap", /* tp_name */
1299 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001300 0, /* tp_itemsize */
1301 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001302 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001303 0, /* tp_print */
1304 0, /* tp_getattr */
1305 0, /* tp_setattr */
1306 0, /* tp_compare */
1307 0, /* tp_repr */
1308 0, /* tp_as_number */
1309 0, /* tp_as_sequence */
1310 0, /* tp_as_mapping */
1311 0, /* tp_hash */
1312 0, /* tp_call */
1313 0, /* tp_str */
1314 PyObject_GenericGetAttr, /* tp_getattro */
1315 0, /* tp_setattro */
1316 0, /* tp_as_buffer */
1317 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1318 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001319 starmap_doc, /* tp_doc */
1320 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321 0, /* tp_clear */
1322 0, /* tp_richcompare */
1323 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001324 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001325 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326 0, /* tp_methods */
1327 0, /* tp_members */
1328 0, /* tp_getset */
1329 0, /* tp_base */
1330 0, /* tp_dict */
1331 0, /* tp_descr_get */
1332 0, /* tp_descr_set */
1333 0, /* tp_dictoffset */
1334 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001335 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001336 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337 PyObject_GC_Del, /* tp_free */
1338};
1339
1340
1341/* imap object ************************************************************/
1342
1343typedef struct {
1344 PyObject_HEAD
1345 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346 PyObject *func;
1347} imapobject;
1348
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001349static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350
1351static PyObject *
1352imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1353{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001354 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001355 imapobject *lz;
1356 int numargs, i;
1357
1358 numargs = PyTuple_Size(args);
1359 if (numargs < 2) {
1360 PyErr_SetString(PyExc_TypeError,
1361 "imap() must have at least two arguments.");
1362 return NULL;
1363 }
1364
1365 iters = PyTuple_New(numargs-1);
1366 if (iters == NULL)
1367 return NULL;
1368
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001369 for (i=1 ; i<numargs ; i++) {
1370 /* Get iterator. */
1371 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1372 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001373 Py_DECREF(iters);
1374 return NULL;
1375 }
1376 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377 }
1378
1379 /* create imapobject structure */
1380 lz = (imapobject *)type->tp_alloc(type, 0);
1381 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001382 Py_DECREF(iters);
1383 return NULL;
1384 }
1385 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001386 func = PyTuple_GET_ITEM(args, 0);
1387 Py_INCREF(func);
1388 lz->func = func;
1389
1390 return (PyObject *)lz;
1391}
1392
1393static void
1394imap_dealloc(imapobject *lz)
1395{
1396 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001397 Py_XDECREF(lz->iters);
1398 Py_XDECREF(lz->func);
1399 lz->ob_type->tp_free(lz);
1400}
1401
1402static int
1403imap_traverse(imapobject *lz, visitproc visit, void *arg)
1404{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001405 Py_VISIT(lz->iters);
1406 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001407 return 0;
1408}
1409
Raymond Hettinger2012f172003-02-07 05:32:58 +00001410/*
1411imap() is an iterator version of __builtins__.map() except that it does
1412not have the None fill-in feature. That was intentionally left out for
1413the following reasons:
1414
1415 1) Itertools are designed to be easily combined and chained together.
1416 Having all tools stop with the shortest input is a unifying principle
1417 that makes it easier to combine finite iterators (supplying data) with
1418 infinite iterators like count() and repeat() (for supplying sequential
1419 or constant arguments to a function).
1420
1421 2) In typical use cases for combining itertools, having one finite data
1422 supplier run out before another is likely to be an error condition which
1423 should not pass silently by automatically supplying None.
1424
1425 3) The use cases for automatic None fill-in are rare -- not many functions
1426 do something useful when a parameter suddenly switches type and becomes
1427 None.
1428
1429 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001430 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001431
1432 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1433*/
1434
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435static PyObject *
1436imap_next(imapobject *lz)
1437{
1438 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001439 PyObject *argtuple;
1440 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001441 int numargs, i;
1442
1443 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001444 argtuple = PyTuple_New(numargs);
1445 if (argtuple == NULL)
1446 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001448 for (i=0 ; i<numargs ; i++) {
1449 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1450 if (val == NULL) {
1451 Py_DECREF(argtuple);
1452 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001454 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001456 if (lz->func == Py_None)
1457 return argtuple;
1458 result = PyObject_Call(lz->func, argtuple, NULL);
1459 Py_DECREF(argtuple);
1460 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461}
1462
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463PyDoc_STRVAR(imap_doc,
1464"imap(func, *iterables) --> imap object\n\
1465\n\
1466Make an iterator that computes the function using arguments from\n\
1467each of the iterables. Like map() except that it returns\n\
1468an iterator instead of a list and that it stops when the shortest\n\
1469iterable is exhausted instead of filling in None for shorter\n\
1470iterables.");
1471
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001472static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473 PyObject_HEAD_INIT(NULL)
1474 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001475 "itertools.imap", /* tp_name */
1476 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477 0, /* tp_itemsize */
1478 /* methods */
1479 (destructor)imap_dealloc, /* tp_dealloc */
1480 0, /* tp_print */
1481 0, /* tp_getattr */
1482 0, /* tp_setattr */
1483 0, /* tp_compare */
1484 0, /* tp_repr */
1485 0, /* tp_as_number */
1486 0, /* tp_as_sequence */
1487 0, /* tp_as_mapping */
1488 0, /* tp_hash */
1489 0, /* tp_call */
1490 0, /* tp_str */
1491 PyObject_GenericGetAttr, /* tp_getattro */
1492 0, /* tp_setattro */
1493 0, /* tp_as_buffer */
1494 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1495 Py_TPFLAGS_BASETYPE, /* tp_flags */
1496 imap_doc, /* tp_doc */
1497 (traverseproc)imap_traverse, /* tp_traverse */
1498 0, /* tp_clear */
1499 0, /* tp_richcompare */
1500 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001501 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001502 (iternextfunc)imap_next, /* tp_iternext */
1503 0, /* tp_methods */
1504 0, /* tp_members */
1505 0, /* tp_getset */
1506 0, /* tp_base */
1507 0, /* tp_dict */
1508 0, /* tp_descr_get */
1509 0, /* tp_descr_set */
1510 0, /* tp_dictoffset */
1511 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001512 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001513 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001514 PyObject_GC_Del, /* tp_free */
1515};
1516
1517
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001518/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001519
1520typedef struct {
1521 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001522 long tuplesize;
1523 long iternum; /* which iterator is active */
1524 PyObject *ittuple; /* tuple of iterators */
1525} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001526
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001527static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528
1529static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001530chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001531{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001532 chainobject *lz;
1533 int tuplesize = PySequence_Length(args);
1534 int i;
1535 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001537 /* obtain iterators */
1538 assert(PyTuple_Check(args));
1539 ittuple = PyTuple_New(tuplesize);
1540 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001542 for (i=0; i < tuplesize; ++i) {
1543 PyObject *item = PyTuple_GET_ITEM(args, i);
1544 PyObject *it = PyObject_GetIter(item);
1545 if (it == NULL) {
1546 if (PyErr_ExceptionMatches(PyExc_TypeError))
1547 PyErr_Format(PyExc_TypeError,
1548 "chain argument #%d must support iteration",
1549 i+1);
1550 Py_DECREF(ittuple);
1551 return NULL;
1552 }
1553 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554 }
1555
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001556 /* create chainobject structure */
1557 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001558 if (lz == NULL) {
1559 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001560 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001561 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001562
1563 lz->ittuple = ittuple;
1564 lz->iternum = 0;
1565 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001566
1567 return (PyObject *)lz;
1568}
1569
1570static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001571chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001572{
1573 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001574 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001575 lz->ob_type->tp_free(lz);
1576}
1577
Raymond Hettinger2012f172003-02-07 05:32:58 +00001578static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001579chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001580{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001581 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001582 return 0;
1583}
1584
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001586chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001587{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001588 PyObject *it;
1589 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001591 while (lz->iternum < lz->tuplesize) {
1592 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1593 item = PyIter_Next(it);
1594 if (item != NULL)
1595 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001596 if (PyErr_Occurred()) {
1597 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1598 PyErr_Clear();
1599 else
1600 return NULL;
1601 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001602 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001603 }
1604 return NULL;
1605}
1606
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001607PyDoc_STRVAR(chain_doc,
1608"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001610Return a chain object whose .next() method returns elements from the\n\
1611first iterable until it is exhausted, then elements from the next\n\
1612iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001613
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001614static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615 PyObject_HEAD_INIT(NULL)
1616 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001617 "itertools.chain", /* tp_name */
1618 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619 0, /* tp_itemsize */
1620 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001621 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001622 0, /* tp_print */
1623 0, /* tp_getattr */
1624 0, /* tp_setattr */
1625 0, /* tp_compare */
1626 0, /* tp_repr */
1627 0, /* tp_as_number */
1628 0, /* tp_as_sequence */
1629 0, /* tp_as_mapping */
1630 0, /* tp_hash */
1631 0, /* tp_call */
1632 0, /* tp_str */
1633 PyObject_GenericGetAttr, /* tp_getattro */
1634 0, /* tp_setattro */
1635 0, /* tp_as_buffer */
1636 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1637 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001638 chain_doc, /* tp_doc */
1639 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001640 0, /* tp_clear */
1641 0, /* tp_richcompare */
1642 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001643 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001644 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645 0, /* tp_methods */
1646 0, /* tp_members */
1647 0, /* tp_getset */
1648 0, /* tp_base */
1649 0, /* tp_dict */
1650 0, /* tp_descr_get */
1651 0, /* tp_descr_set */
1652 0, /* tp_dictoffset */
1653 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001654 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001655 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001656 PyObject_GC_Del, /* tp_free */
1657};
1658
1659
1660/* ifilter object ************************************************************/
1661
1662typedef struct {
1663 PyObject_HEAD
1664 PyObject *func;
1665 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666} ifilterobject;
1667
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001668static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001669
1670static PyObject *
1671ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1672{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001673 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674 PyObject *it;
1675 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676
Raymond Hettinger60eca932003-02-09 06:40:58 +00001677 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001678 return NULL;
1679
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680 /* Get iterator. */
1681 it = PyObject_GetIter(seq);
1682 if (it == NULL)
1683 return NULL;
1684
1685 /* create ifilterobject structure */
1686 lz = (ifilterobject *)type->tp_alloc(type, 0);
1687 if (lz == NULL) {
1688 Py_DECREF(it);
1689 return NULL;
1690 }
1691 Py_INCREF(func);
1692 lz->func = func;
1693 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694
1695 return (PyObject *)lz;
1696}
1697
1698static void
1699ifilter_dealloc(ifilterobject *lz)
1700{
1701 PyObject_GC_UnTrack(lz);
1702 Py_XDECREF(lz->func);
1703 Py_XDECREF(lz->it);
1704 lz->ob_type->tp_free(lz);
1705}
1706
1707static int
1708ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1709{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001710 Py_VISIT(lz->it);
1711 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001712 return 0;
1713}
1714
1715static PyObject *
1716ifilter_next(ifilterobject *lz)
1717{
1718 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001719 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001720 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001721 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001722
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001723 assert(PyIter_Check(it));
1724 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001725 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001726 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001727 if (item == NULL)
1728 return NULL;
1729
1730 if (lz->func == Py_None) {
1731 ok = PyObject_IsTrue(item);
1732 } else {
1733 PyObject *good;
1734 good = PyObject_CallFunctionObjArgs(lz->func,
1735 item, NULL);
1736 if (good == NULL) {
1737 Py_DECREF(item);
1738 return NULL;
1739 }
1740 ok = PyObject_IsTrue(good);
1741 Py_DECREF(good);
1742 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001743 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001744 return item;
1745 Py_DECREF(item);
1746 }
1747}
1748
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001750"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001751\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001752Return those items of sequence for which function(item) is true.\n\
1753If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001754
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001755static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001756 PyObject_HEAD_INIT(NULL)
1757 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001758 "itertools.ifilter", /* tp_name */
1759 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001760 0, /* tp_itemsize */
1761 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001762 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001763 0, /* tp_print */
1764 0, /* tp_getattr */
1765 0, /* tp_setattr */
1766 0, /* tp_compare */
1767 0, /* tp_repr */
1768 0, /* tp_as_number */
1769 0, /* tp_as_sequence */
1770 0, /* tp_as_mapping */
1771 0, /* tp_hash */
1772 0, /* tp_call */
1773 0, /* tp_str */
1774 PyObject_GenericGetAttr, /* tp_getattro */
1775 0, /* tp_setattro */
1776 0, /* tp_as_buffer */
1777 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1778 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001779 ifilter_doc, /* tp_doc */
1780 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001781 0, /* tp_clear */
1782 0, /* tp_richcompare */
1783 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001784 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001785 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001786 0, /* tp_methods */
1787 0, /* tp_members */
1788 0, /* tp_getset */
1789 0, /* tp_base */
1790 0, /* tp_dict */
1791 0, /* tp_descr_get */
1792 0, /* tp_descr_set */
1793 0, /* tp_dictoffset */
1794 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001795 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001796 ifilter_new, /* tp_new */
1797 PyObject_GC_Del, /* tp_free */
1798};
1799
1800
1801/* ifilterfalse object ************************************************************/
1802
1803typedef struct {
1804 PyObject_HEAD
1805 PyObject *func;
1806 PyObject *it;
1807} ifilterfalseobject;
1808
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001809static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001810
1811static PyObject *
1812ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1813{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001814 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001815 PyObject *it;
1816 ifilterfalseobject *lz;
1817
1818 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1819 return NULL;
1820
1821 /* Get iterator. */
1822 it = PyObject_GetIter(seq);
1823 if (it == NULL)
1824 return NULL;
1825
1826 /* create ifilterfalseobject structure */
1827 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1828 if (lz == NULL) {
1829 Py_DECREF(it);
1830 return NULL;
1831 }
1832 Py_INCREF(func);
1833 lz->func = func;
1834 lz->it = it;
1835
1836 return (PyObject *)lz;
1837}
1838
1839static void
1840ifilterfalse_dealloc(ifilterfalseobject *lz)
1841{
1842 PyObject_GC_UnTrack(lz);
1843 Py_XDECREF(lz->func);
1844 Py_XDECREF(lz->it);
1845 lz->ob_type->tp_free(lz);
1846}
1847
1848static int
1849ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1850{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001851 Py_VISIT(lz->it);
1852 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001853 return 0;
1854}
1855
1856static PyObject *
1857ifilterfalse_next(ifilterfalseobject *lz)
1858{
1859 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001860 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001861 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001862 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001863
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001864 assert(PyIter_Check(it));
1865 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001866 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001867 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001868 if (item == NULL)
1869 return NULL;
1870
1871 if (lz->func == Py_None) {
1872 ok = PyObject_IsTrue(item);
1873 } else {
1874 PyObject *good;
1875 good = PyObject_CallFunctionObjArgs(lz->func,
1876 item, NULL);
1877 if (good == NULL) {
1878 Py_DECREF(item);
1879 return NULL;
1880 }
1881 ok = PyObject_IsTrue(good);
1882 Py_DECREF(good);
1883 }
1884 if (!ok)
1885 return item;
1886 Py_DECREF(item);
1887 }
1888}
1889
Raymond Hettinger60eca932003-02-09 06:40:58 +00001890PyDoc_STRVAR(ifilterfalse_doc,
1891"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1892\n\
1893Return those items of sequence for which function(item) is false.\n\
1894If function is None, return the items that are false.");
1895
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001896static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001897 PyObject_HEAD_INIT(NULL)
1898 0, /* ob_size */
1899 "itertools.ifilterfalse", /* tp_name */
1900 sizeof(ifilterfalseobject), /* tp_basicsize */
1901 0, /* tp_itemsize */
1902 /* methods */
1903 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1904 0, /* tp_print */
1905 0, /* tp_getattr */
1906 0, /* tp_setattr */
1907 0, /* tp_compare */
1908 0, /* tp_repr */
1909 0, /* tp_as_number */
1910 0, /* tp_as_sequence */
1911 0, /* tp_as_mapping */
1912 0, /* tp_hash */
1913 0, /* tp_call */
1914 0, /* tp_str */
1915 PyObject_GenericGetAttr, /* tp_getattro */
1916 0, /* tp_setattro */
1917 0, /* tp_as_buffer */
1918 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1919 Py_TPFLAGS_BASETYPE, /* tp_flags */
1920 ifilterfalse_doc, /* tp_doc */
1921 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1922 0, /* tp_clear */
1923 0, /* tp_richcompare */
1924 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001925 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001926 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1927 0, /* tp_methods */
1928 0, /* tp_members */
1929 0, /* tp_getset */
1930 0, /* tp_base */
1931 0, /* tp_dict */
1932 0, /* tp_descr_get */
1933 0, /* tp_descr_set */
1934 0, /* tp_dictoffset */
1935 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001936 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001937 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001938 PyObject_GC_Del, /* tp_free */
1939};
1940
1941
1942/* count object ************************************************************/
1943
1944typedef struct {
1945 PyObject_HEAD
1946 long cnt;
1947} countobject;
1948
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001949static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001950
1951static PyObject *
1952count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1953{
1954 countobject *lz;
1955 long cnt = 0;
1956
1957 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1958 return NULL;
1959
1960 /* create countobject structure */
1961 lz = (countobject *)PyObject_New(countobject, &count_type);
1962 if (lz == NULL)
1963 return NULL;
1964 lz->cnt = cnt;
1965
1966 return (PyObject *)lz;
1967}
1968
1969static PyObject *
1970count_next(countobject *lz)
1971{
1972 return PyInt_FromLong(lz->cnt++);
1973}
1974
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001975static PyObject *
1976count_repr(countobject *lz)
1977{
Brett Cannon00468392004-04-13 02:43:53 +00001978 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001979}
1980
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001981PyDoc_STRVAR(count_doc,
1982"count([firstval]) --> count object\n\
1983\n\
1984Return a count object whose .next() method returns consecutive\n\
1985integers starting from zero or, if specified, from firstval.");
1986
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001987static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001988 PyObject_HEAD_INIT(NULL)
1989 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001990 "itertools.count", /* tp_name */
1991 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001992 0, /* tp_itemsize */
1993 /* methods */
1994 (destructor)PyObject_Del, /* tp_dealloc */
1995 0, /* tp_print */
1996 0, /* tp_getattr */
1997 0, /* tp_setattr */
1998 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001999 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002000 0, /* tp_as_number */
2001 0, /* tp_as_sequence */
2002 0, /* tp_as_mapping */
2003 0, /* tp_hash */
2004 0, /* tp_call */
2005 0, /* tp_str */
2006 PyObject_GenericGetAttr, /* tp_getattro */
2007 0, /* tp_setattro */
2008 0, /* tp_as_buffer */
2009 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002010 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002011 0, /* tp_traverse */
2012 0, /* tp_clear */
2013 0, /* tp_richcompare */
2014 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002015 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002016 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002017 0, /* tp_methods */
2018 0, /* tp_members */
2019 0, /* tp_getset */
2020 0, /* tp_base */
2021 0, /* tp_dict */
2022 0, /* tp_descr_get */
2023 0, /* tp_descr_set */
2024 0, /* tp_dictoffset */
2025 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002026 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002027 count_new, /* tp_new */
2028};
2029
2030
2031/* izip object ************************************************************/
2032
2033#include "Python.h"
2034
2035typedef struct {
2036 PyObject_HEAD
2037 long tuplesize;
2038 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002039 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002040} izipobject;
2041
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002042static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002043
2044static PyObject *
2045izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2046{
2047 izipobject *lz;
2048 int i;
2049 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002050 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002051 int tuplesize = PySequence_Length(args);
2052
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002053 /* args must be a tuple */
2054 assert(PyTuple_Check(args));
2055
2056 /* obtain iterators */
2057 ittuple = PyTuple_New(tuplesize);
2058 if(ittuple == NULL)
2059 return NULL;
2060 for (i=0; i < tuplesize; ++i) {
2061 PyObject *item = PyTuple_GET_ITEM(args, i);
2062 PyObject *it = PyObject_GetIter(item);
2063 if (it == NULL) {
2064 if (PyErr_ExceptionMatches(PyExc_TypeError))
2065 PyErr_Format(PyExc_TypeError,
2066 "izip argument #%d must support iteration",
2067 i+1);
2068 Py_DECREF(ittuple);
2069 return NULL;
2070 }
2071 PyTuple_SET_ITEM(ittuple, i, it);
2072 }
2073
Raymond Hettinger2012f172003-02-07 05:32:58 +00002074 /* create a result holder */
2075 result = PyTuple_New(tuplesize);
2076 if (result == NULL) {
2077 Py_DECREF(ittuple);
2078 return NULL;
2079 }
2080 for (i=0 ; i < tuplesize ; i++) {
2081 Py_INCREF(Py_None);
2082 PyTuple_SET_ITEM(result, i, Py_None);
2083 }
2084
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002085 /* create izipobject structure */
2086 lz = (izipobject *)type->tp_alloc(type, 0);
2087 if (lz == NULL) {
2088 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002089 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002090 return NULL;
2091 }
2092 lz->ittuple = ittuple;
2093 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002094 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002095
2096 return (PyObject *)lz;
2097}
2098
2099static void
2100izip_dealloc(izipobject *lz)
2101{
2102 PyObject_GC_UnTrack(lz);
2103 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002104 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002105 lz->ob_type->tp_free(lz);
2106}
2107
2108static int
2109izip_traverse(izipobject *lz, visitproc visit, void *arg)
2110{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002111 Py_VISIT(lz->ittuple);
2112 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002113 return 0;
2114}
2115
2116static PyObject *
2117izip_next(izipobject *lz)
2118{
2119 int i;
2120 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002121 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002122 PyObject *it;
2123 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002124 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002125
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002126 if (tuplesize == 0)
2127 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002128 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002129 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002130 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002131 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002132 assert(PyIter_Check(it));
2133 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002134 if (item == NULL) {
2135 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002136 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002137 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002138 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002139 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002140 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002141 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002142 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002143 result = PyTuple_New(tuplesize);
2144 if (result == NULL)
2145 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002146 for (i=0 ; i < tuplesize ; i++) {
2147 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002148 assert(PyIter_Check(it));
2149 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002150 if (item == NULL) {
2151 Py_DECREF(result);
2152 return NULL;
2153 }
2154 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002155 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002156 }
2157 return result;
2158}
2159
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002160PyDoc_STRVAR(izip_doc,
2161"izip(iter1 [,iter2 [...]]) --> izip object\n\
2162\n\
2163Return a izip object whose .next() method returns a tuple where\n\
2164the i-th element comes from the i-th iterable argument. The .next()\n\
2165method continues until the shortest iterable in the argument sequence\n\
2166is exhausted and then it raises StopIteration. Works like the zip()\n\
2167function but consumes less memory by returning an iterator instead of\n\
2168a list.");
2169
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002170static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002171 PyObject_HEAD_INIT(NULL)
2172 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002173 "itertools.izip", /* tp_name */
2174 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002175 0, /* tp_itemsize */
2176 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002177 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002178 0, /* tp_print */
2179 0, /* tp_getattr */
2180 0, /* tp_setattr */
2181 0, /* tp_compare */
2182 0, /* tp_repr */
2183 0, /* tp_as_number */
2184 0, /* tp_as_sequence */
2185 0, /* tp_as_mapping */
2186 0, /* tp_hash */
2187 0, /* tp_call */
2188 0, /* tp_str */
2189 PyObject_GenericGetAttr, /* tp_getattro */
2190 0, /* tp_setattro */
2191 0, /* tp_as_buffer */
2192 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2193 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002194 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002195 (traverseproc)izip_traverse, /* tp_traverse */
2196 0, /* tp_clear */
2197 0, /* tp_richcompare */
2198 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002199 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002200 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002201 0, /* tp_methods */
2202 0, /* tp_members */
2203 0, /* tp_getset */
2204 0, /* tp_base */
2205 0, /* tp_dict */
2206 0, /* tp_descr_get */
2207 0, /* tp_descr_set */
2208 0, /* tp_dictoffset */
2209 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002210 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002211 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002212 PyObject_GC_Del, /* tp_free */
2213};
2214
2215
2216/* repeat object ************************************************************/
2217
2218typedef struct {
2219 PyObject_HEAD
2220 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002221 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002222} repeatobject;
2223
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002224static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002225
2226static PyObject *
2227repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2228{
2229 repeatobject *ro;
2230 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002231 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002232
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002233 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002234 return NULL;
2235
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002236 if (PyTuple_Size(args) == 2 && cnt < 0)
2237 cnt = 0;
2238
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002239 ro = (repeatobject *)type->tp_alloc(type, 0);
2240 if (ro == NULL)
2241 return NULL;
2242 Py_INCREF(element);
2243 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002244 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002245 return (PyObject *)ro;
2246}
2247
2248static void
2249repeat_dealloc(repeatobject *ro)
2250{
2251 PyObject_GC_UnTrack(ro);
2252 Py_XDECREF(ro->element);
2253 ro->ob_type->tp_free(ro);
2254}
2255
2256static int
2257repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2258{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002259 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002260 return 0;
2261}
2262
2263static PyObject *
2264repeat_next(repeatobject *ro)
2265{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002266 if (ro->cnt == 0)
2267 return NULL;
2268 if (ro->cnt > 0)
2269 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002270 Py_INCREF(ro->element);
2271 return ro->element;
2272}
2273
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002274static PyObject *
2275repeat_repr(repeatobject *ro)
2276{
2277 PyObject *result, *objrepr;
2278
2279 objrepr = PyObject_Repr(ro->element);
2280 if (objrepr == NULL)
2281 return NULL;
2282
2283 if (ro->cnt == -1)
2284 result = PyString_FromFormat("repeat(%s)",
2285 PyString_AS_STRING(objrepr));
2286 else
Brett Cannon00468392004-04-13 02:43:53 +00002287 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002288 PyString_AS_STRING(objrepr), ro->cnt);
2289 Py_DECREF(objrepr);
2290 return result;
2291}
2292
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002293static int
2294repeat_len(repeatobject *ro)
2295{
2296 if (ro->cnt == -1)
2297 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2298 return (int)(ro->cnt);
2299}
2300
2301static PySequenceMethods repeat_as_sequence = {
2302 (inquiry)repeat_len, /* sq_length */
2303 0, /* sq_concat */
2304};
2305
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002306PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002307"repeat(element [,times]) -> create an iterator which returns the element\n\
2308for the specified number of times. If not specified, returns the element\n\
2309endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002310
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002311static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002312 PyObject_HEAD_INIT(NULL)
2313 0, /* ob_size */
2314 "itertools.repeat", /* tp_name */
2315 sizeof(repeatobject), /* tp_basicsize */
2316 0, /* tp_itemsize */
2317 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002318 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002319 0, /* tp_print */
2320 0, /* tp_getattr */
2321 0, /* tp_setattr */
2322 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002323 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002324 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002325 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002326 0, /* tp_as_mapping */
2327 0, /* tp_hash */
2328 0, /* tp_call */
2329 0, /* tp_str */
2330 PyObject_GenericGetAttr, /* tp_getattro */
2331 0, /* tp_setattro */
2332 0, /* tp_as_buffer */
2333 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2334 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002335 repeat_doc, /* tp_doc */
2336 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002337 0, /* tp_clear */
2338 0, /* tp_richcompare */
2339 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002340 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002341 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002342 0, /* tp_methods */
2343 0, /* tp_members */
2344 0, /* tp_getset */
2345 0, /* tp_base */
2346 0, /* tp_dict */
2347 0, /* tp_descr_get */
2348 0, /* tp_descr_set */
2349 0, /* tp_dictoffset */
2350 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002351 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002352 repeat_new, /* tp_new */
2353 PyObject_GC_Del, /* tp_free */
2354};
2355
2356
2357/* module level code ********************************************************/
2358
2359PyDoc_STRVAR(module_doc,
2360"Functional tools for creating and using iterators.\n\
2361\n\
2362Infinite iterators:\n\
2363count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002364cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002365repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002366\n\
2367Iterators terminating on the shortest input sequence:\n\
2368izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002369ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2370ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002371islice(seq, [start,] stop [, step]) --> elements from\n\
2372 seq[start:stop:step]\n\
2373imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2374starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002375tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002376chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002377takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2378dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002379groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002380");
2381
2382
Raymond Hettingerad983e72003-11-12 14:32:26 +00002383static PyMethodDef module_methods[] = {
2384 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2385 {NULL, NULL} /* sentinel */
2386};
2387
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002388PyMODINIT_FUNC
2389inititertools(void)
2390{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002391 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002392 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002393 char *name;
2394 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002395 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002396 &dropwhile_type,
2397 &takewhile_type,
2398 &islice_type,
2399 &starmap_type,
2400 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002401 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002402 &ifilter_type,
2403 &ifilterfalse_type,
2404 &count_type,
2405 &izip_type,
2406 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002407 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002408 NULL
2409 };
2410
Skip Montanarof3938fd2004-02-10 20:27:40 +00002411 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002412 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002413
Raymond Hettinger60eca932003-02-09 06:40:58 +00002414 for (i=0 ; typelist[i] != NULL ; i++) {
2415 if (PyType_Ready(typelist[i]) < 0)
2416 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002417 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002418 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002419 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002420 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002421 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002422
2423 if (PyType_Ready(&teedataobject_type) < 0)
2424 return;
2425 if (PyType_Ready(&tee_type) < 0)
2426 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002427 if (PyType_Ready(&_grouper_type) < 0)
2428 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002429}