blob: 31ba13a3ad6dbab91290c025469a34e75fb3c8e3 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
5/* Itertools module written and maintained
6 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
15 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
21} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
29 static char *kwargs[] = {"iterable", "key", NULL};
30 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
32
33 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
51}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
56 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
62 gbo->ob_type->tp_free(gbo);
63}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000073 return 0;
74}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
81 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
89
90 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
97
98 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
101
102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
113
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000114 tmp = gbo->currkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000115 gbo->currkey = newkey;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000116 Py_XDECREF(tmp);
117
118 tmp = gbo->currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000119 gbo->currvalue = newvalue;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000120 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121 }
122
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000123 Py_INCREF(gbo->currkey);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
131
132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
142 PyObject_HEAD_INIT(NULL)
143 0, /* ob_size */
144 "itertools.groupby", /* tp_name */
145 sizeof(groupbyobject), /* tp_basicsize */
146 0, /* tp_itemsize */
147 /* methods */
148 (destructor)groupby_dealloc, /* tp_dealloc */
149 0, /* tp_print */
150 0, /* tp_getattr */
151 0, /* tp_setattr */
152 0, /* tp_compare */
153 0, /* tp_repr */
154 0, /* tp_as_number */
155 0, /* tp_as_sequence */
156 0, /* tp_as_mapping */
157 0, /* tp_hash */
158 0, /* tp_call */
159 0, /* tp_str */
160 PyObject_GenericGetAttr, /* tp_getattro */
161 0, /* tp_setattro */
162 0, /* tp_as_buffer */
163 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
164 Py_TPFLAGS_BASETYPE, /* tp_flags */
165 groupby_doc, /* tp_doc */
166 (traverseproc)groupby_traverse, /* tp_traverse */
167 0, /* tp_clear */
168 0, /* tp_richcompare */
169 0, /* tp_weaklistoffset */
170 PyObject_SelfIter, /* tp_iter */
171 (iternextfunc)groupby_next, /* tp_iternext */
172 0, /* tp_methods */
173 0, /* tp_members */
174 0, /* tp_getset */
175 0, /* tp_base */
176 0, /* tp_dict */
177 0, /* tp_descr_get */
178 0, /* tp_descr_set */
179 0, /* tp_dictoffset */
180 0, /* tp_init */
181 0, /* tp_alloc */
182 groupby_new, /* tp_new */
183 PyObject_GC_Del, /* tp_free */
184};
185
186
187/* _grouper object (internal) ************************************************/
188
189typedef struct {
190 PyObject_HEAD
191 PyObject *parent;
192 PyObject *tgtkey;
193} _grouperobject;
194
195static PyTypeObject _grouper_type;
196
197static PyObject *
198_grouper_create(groupbyobject *parent, PyObject *tgtkey)
199{
200 _grouperobject *igo;
201
202 igo = PyObject_New(_grouperobject, &_grouper_type);
203 if (igo == NULL)
204 return NULL;
205 igo->parent = (PyObject *)parent;
206 Py_INCREF(parent);
207 igo->tgtkey = tgtkey;
208 Py_INCREF(tgtkey);
209
210 return (PyObject *)igo;
211}
212
213static void
214_grouper_dealloc(_grouperobject *igo)
215{
216 Py_DECREF(igo->parent);
217 Py_DECREF(igo->tgtkey);
218 PyObject_Del(igo);
219}
220
221static PyObject *
222_grouper_next(_grouperobject *igo)
223{
224 groupbyobject *gbo = (groupbyobject *)igo->parent;
225 PyObject *newvalue, *newkey, *r;
226 int rcmp;
227
228 if (gbo->currvalue == NULL) {
229 newvalue = PyIter_Next(gbo->it);
230 if (newvalue == NULL)
231 return NULL;
232
233 if (gbo->keyfunc == Py_None) {
234 newkey = newvalue;
235 Py_INCREF(newvalue);
236 } else {
237 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
238 newvalue, NULL);
239 if (newkey == NULL) {
240 Py_DECREF(newvalue);
241 return NULL;
242 }
243 }
244
245 assert(gbo->currkey == NULL);
246 gbo->currkey = newkey;
247 gbo->currvalue = newvalue;
248 }
249
250 assert(gbo->currkey != NULL);
251 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
252 if (rcmp <= 0)
253 /* got any error or current group is end */
254 return NULL;
255
256 r = gbo->currvalue;
257 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000258 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000259
260 return r;
261}
262
263static PyTypeObject _grouper_type = {
264 PyObject_HEAD_INIT(NULL)
265 0, /* ob_size */
266 "itertools._grouper", /* tp_name */
267 sizeof(_grouperobject), /* tp_basicsize */
268 0, /* tp_itemsize */
269 /* methods */
270 (destructor)_grouper_dealloc, /* tp_dealloc */
271 0, /* tp_print */
272 0, /* tp_getattr */
273 0, /* tp_setattr */
274 0, /* tp_compare */
275 0, /* tp_repr */
276 0, /* tp_as_number */
277 0, /* tp_as_sequence */
278 0, /* tp_as_mapping */
279 0, /* tp_hash */
280 0, /* tp_call */
281 0, /* tp_str */
282 PyObject_GenericGetAttr, /* tp_getattro */
283 0, /* tp_setattro */
284 0, /* tp_as_buffer */
285 Py_TPFLAGS_DEFAULT, /* tp_flags */
286 0, /* tp_doc */
287 0, /* tp_traverse */
288 0, /* tp_clear */
289 0, /* tp_richcompare */
290 0, /* tp_weaklistoffset */
291 PyObject_SelfIter, /* tp_iter */
292 (iternextfunc)_grouper_next, /* tp_iternext */
293 0, /* tp_methods */
294 0, /* tp_members */
295 0, /* tp_getset */
296 0, /* tp_base */
297 0, /* tp_dict */
298 0, /* tp_descr_get */
299 0, /* tp_descr_set */
300 0, /* tp_dictoffset */
301 0, /* tp_init */
302 0, /* tp_alloc */
303 0, /* tp_new */
304 PyObject_Del, /* tp_free */
305};
306
307
308
Raymond Hettingerad983e72003-11-12 14:32:26 +0000309/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000310
Raymond Hettingerad983e72003-11-12 14:32:26 +0000311/* The teedataobject pre-allocates space for LINKCELLS number of objects.
312 To help the object fit neatly inside cache lines (space for 16 to 32
313 pointers), the value should be a multiple of 16 minus space for
314 the other structure members including PyHEAD overhead. The larger the
315 value, the less memory overhead per object and the less time spent
316 allocating/deallocating new links. The smaller the number, the less
317 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000320
321typedef struct {
322 PyObject_HEAD
323 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000324 int numread;
325 PyObject *nextlink;
326 PyObject *(values[LINKCELLS]);
327} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
329typedef struct {
330 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000331 teedataobject *dataobj;
332 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000333 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335
Raymond Hettingerad983e72003-11-12 14:32:26 +0000336static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000337
338static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000340{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000341 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000342
Raymond Hettingerad983e72003-11-12 14:32:26 +0000343 tdo = PyObject_New(teedataobject, &teedataobject_type);
344 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000345 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000346
347 tdo->numread = 0;
348 tdo->nextlink = NULL;
349 Py_INCREF(it);
350 tdo->it = it;
351 return (PyObject *)tdo;
352}
353
354static PyObject *
355teedataobject_jumplink(teedataobject *tdo)
356{
357 if (tdo->nextlink == NULL)
358 tdo->nextlink = teedataobject_new(tdo->it);
359 Py_INCREF(tdo->nextlink);
360 return tdo->nextlink;
361}
362
363static PyObject *
364teedataobject_getitem(teedataobject *tdo, int i)
365{
366 PyObject *value;
367
368 assert(i < LINKCELLS);
369 if (i < tdo->numread)
370 value = tdo->values[i];
371 else {
372 /* this is the lead iterator, so fetch more data */
373 assert(i == tdo->numread);
374 value = PyIter_Next(tdo->it);
375 if (value == NULL)
376 return NULL;
377 tdo->numread++;
378 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000379 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000380 Py_INCREF(value);
381 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000382}
383
384static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000385teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000386{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000387 int i;
Raymond Hettinger45143692003-10-25 06:37:47 +0000388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 for (i=0 ; i<tdo->numread ; i++)
390 Py_DECREF(tdo->values[i]);
391 Py_XDECREF(tdo->it);
392 Py_XDECREF(tdo->nextlink);
393 PyObject_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000394}
395
Raymond Hettingerad983e72003-11-12 14:32:26 +0000396PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000397
Raymond Hettingerad983e72003-11-12 14:32:26 +0000398static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000399 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000401 "itertools.tee_dataobject", /* tp_name */
402 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000403 0, /* tp_itemsize */
404 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000405 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000406 0, /* tp_print */
407 0, /* tp_getattr */
408 0, /* tp_setattr */
409 0, /* tp_compare */
410 0, /* tp_repr */
411 0, /* tp_as_number */
412 0, /* tp_as_sequence */
413 0, /* tp_as_mapping */
414 0, /* tp_hash */
415 0, /* tp_call */
416 0, /* tp_str */
417 PyObject_GenericGetAttr, /* tp_getattro */
418 0, /* tp_setattro */
419 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000420 Py_TPFLAGS_DEFAULT, /* tp_flags */
421 teedataobject_doc, /* tp_doc */
422 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000423};
424
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000425
426static PyTypeObject tee_type;
427
428static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000429tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000430{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000431 PyObject *value, *link;
432
433 if (to->index >= LINKCELLS) {
434 link = teedataobject_jumplink(to->dataobj);
435 Py_XDECREF(to->dataobj);
436 to->dataobj = (teedataobject *)link;
437 to->index = 0;
438 }
439 value = teedataobject_getitem(to->dataobj, to->index);
440 if (value == NULL)
441 return NULL;
442 to->index++;
443 return value;
444}
445
446static PyObject *
447tee_copy(teeobject *to)
448{
449 teeobject *newto;
450
451 newto = PyObject_New(teeobject, &tee_type);
452 if (newto == NULL)
453 return NULL;
454 Py_INCREF(to->dataobj);
455 newto->dataobj = to->dataobj;
456 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000457 newto->weakreflist = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000458 return (PyObject *)newto;
459}
460
461PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
462
463static PyObject *
464tee_fromiterable(PyObject *iterable)
465{
466 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000467 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000468
469 it = PyObject_GetIter(iterable);
470 if (it == NULL)
471 return NULL;
472 if (PyObject_TypeCheck(it, &tee_type)) {
473 to = (teeobject *)tee_copy((teeobject *)it);
474 goto done;
475 }
476
477 to = PyObject_New(teeobject, &tee_type);
478 if (to == NULL)
479 goto done;
480 to->dataobj = (teedataobject *)teedataobject_new(it);
481 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000482 to->weakreflist = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000483done:
484 Py_XDECREF(it);
485 return (PyObject *)to;
486}
487
488static PyObject *
489tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
490{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000491 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000492
493 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
494 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000495 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000496}
497
498static void
499tee_dealloc(teeobject *to)
500{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000501 if (to->weakreflist != NULL)
502 PyObject_ClearWeakRefs((PyObject *) to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000503 Py_XDECREF(to->dataobj);
504 PyObject_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000505}
506
Raymond Hettingerad983e72003-11-12 14:32:26 +0000507PyDoc_STRVAR(teeobject_doc,
508"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000509
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510static PyMethodDef tee_methods[] = {
511 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
512 {NULL, NULL} /* sentinel */
513};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000514
515static PyTypeObject tee_type = {
516 PyObject_HEAD_INIT(NULL)
517 0, /* ob_size */
518 "itertools.tee", /* tp_name */
519 sizeof(teeobject), /* tp_basicsize */
520 0, /* tp_itemsize */
521 /* methods */
522 (destructor)tee_dealloc, /* tp_dealloc */
523 0, /* tp_print */
524 0, /* tp_getattr */
525 0, /* tp_setattr */
526 0, /* tp_compare */
527 0, /* tp_repr */
528 0, /* tp_as_number */
529 0, /* tp_as_sequence */
530 0, /* tp_as_mapping */
531 0, /* tp_hash */
532 0, /* tp_call */
533 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000534 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000535 0, /* tp_setattro */
536 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000537 Py_TPFLAGS_DEFAULT, /* tp_flags */
538 teeobject_doc, /* tp_doc */
539 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000540 0, /* tp_clear */
541 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000542 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000543 PyObject_SelfIter, /* tp_iter */
544 (iternextfunc)tee_next, /* tp_iternext */
545 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000546 0, /* tp_members */
547 0, /* tp_getset */
548 0, /* tp_base */
549 0, /* tp_dict */
550 0, /* tp_descr_get */
551 0, /* tp_descr_set */
552 0, /* tp_dictoffset */
553 0, /* tp_init */
554 0, /* tp_alloc */
555 tee_new, /* tp_new */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000556 PyObject_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000557};
558
Raymond Hettingerad983e72003-11-12 14:32:26 +0000559static PyObject *
560tee(PyObject *self, PyObject *args)
561{
562 int i, n=2;
563 PyObject *it, *iterable, *copyable, *result;
564
565 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
566 return NULL;
567 result = PyTuple_New(n);
568 if (result == NULL)
569 return NULL;
570 if (n == 0)
571 return result;
572 it = PyObject_GetIter(iterable);
573 if (it == NULL) {
574 Py_DECREF(result);
575 return NULL;
576 }
577 if (!PyObject_HasAttrString(it, "__copy__")) {
578 copyable = tee_fromiterable(it);
579 Py_DECREF(it);
580 if (copyable == NULL) {
581 Py_DECREF(result);
582 return NULL;
583 }
584 } else
585 copyable = it;
586 PyTuple_SET_ITEM(result, 0, copyable);
587 for (i=1 ; i<n ; i++) {
588 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
589 if (copyable == NULL) {
590 Py_DECREF(result);
591 return NULL;
592 }
593 PyTuple_SET_ITEM(result, i, copyable);
594 }
595 return result;
596}
597
598PyDoc_STRVAR(tee_doc,
599"tee(iterable, n=2) --> tuple of n independent iterators.");
600
601
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000602/* cycle object **********************************************************/
603
604typedef struct {
605 PyObject_HEAD
606 PyObject *it;
607 PyObject *saved;
608 int firstpass;
609} cycleobject;
610
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000611static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000612
613static PyObject *
614cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
615{
616 PyObject *it;
617 PyObject *iterable;
618 PyObject *saved;
619 cycleobject *lz;
620
621 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
622 return NULL;
623
624 /* Get iterator. */
625 it = PyObject_GetIter(iterable);
626 if (it == NULL)
627 return NULL;
628
629 saved = PyList_New(0);
630 if (saved == NULL) {
631 Py_DECREF(it);
632 return NULL;
633 }
634
635 /* create cycleobject structure */
636 lz = (cycleobject *)type->tp_alloc(type, 0);
637 if (lz == NULL) {
638 Py_DECREF(it);
639 Py_DECREF(saved);
640 return NULL;
641 }
642 lz->it = it;
643 lz->saved = saved;
644 lz->firstpass = 0;
645
646 return (PyObject *)lz;
647}
648
649static void
650cycle_dealloc(cycleobject *lz)
651{
652 PyObject_GC_UnTrack(lz);
653 Py_XDECREF(lz->saved);
654 Py_XDECREF(lz->it);
655 lz->ob_type->tp_free(lz);
656}
657
658static int
659cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
660{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000661 Py_VISIT(lz->it);
662 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000663 return 0;
664}
665
666static PyObject *
667cycle_next(cycleobject *lz)
668{
669 PyObject *item;
670 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000671 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000672
673 while (1) {
674 item = PyIter_Next(lz->it);
675 if (item != NULL) {
676 if (!lz->firstpass)
677 PyList_Append(lz->saved, item);
678 return item;
679 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000680 if (PyErr_Occurred()) {
681 if (PyErr_ExceptionMatches(PyExc_StopIteration))
682 PyErr_Clear();
683 else
684 return NULL;
685 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000686 if (PyList_Size(lz->saved) == 0)
687 return NULL;
688 it = PyObject_GetIter(lz->saved);
689 if (it == NULL)
690 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000691 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000692 lz->it = it;
693 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000694 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000695 }
696}
697
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000698PyDoc_STRVAR(cycle_doc,
699"cycle(iterable) --> cycle object\n\
700\n\
701Return elements from the iterable until it is exhausted.\n\
702Then repeat the sequence indefinitely.");
703
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000704static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000705 PyObject_HEAD_INIT(NULL)
706 0, /* ob_size */
707 "itertools.cycle", /* tp_name */
708 sizeof(cycleobject), /* tp_basicsize */
709 0, /* tp_itemsize */
710 /* methods */
711 (destructor)cycle_dealloc, /* tp_dealloc */
712 0, /* tp_print */
713 0, /* tp_getattr */
714 0, /* tp_setattr */
715 0, /* tp_compare */
716 0, /* tp_repr */
717 0, /* tp_as_number */
718 0, /* tp_as_sequence */
719 0, /* tp_as_mapping */
720 0, /* tp_hash */
721 0, /* tp_call */
722 0, /* tp_str */
723 PyObject_GenericGetAttr, /* tp_getattro */
724 0, /* tp_setattro */
725 0, /* tp_as_buffer */
726 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
727 Py_TPFLAGS_BASETYPE, /* tp_flags */
728 cycle_doc, /* tp_doc */
729 (traverseproc)cycle_traverse, /* tp_traverse */
730 0, /* tp_clear */
731 0, /* tp_richcompare */
732 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000733 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000734 (iternextfunc)cycle_next, /* tp_iternext */
735 0, /* tp_methods */
736 0, /* tp_members */
737 0, /* tp_getset */
738 0, /* tp_base */
739 0, /* tp_dict */
740 0, /* tp_descr_get */
741 0, /* tp_descr_set */
742 0, /* tp_dictoffset */
743 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000744 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000745 cycle_new, /* tp_new */
746 PyObject_GC_Del, /* tp_free */
747};
748
749
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000750/* dropwhile object **********************************************************/
751
752typedef struct {
753 PyObject_HEAD
754 PyObject *func;
755 PyObject *it;
756 long start;
757} dropwhileobject;
758
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000759static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000760
761static PyObject *
762dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
763{
764 PyObject *func, *seq;
765 PyObject *it;
766 dropwhileobject *lz;
767
768 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
769 return NULL;
770
771 /* Get iterator. */
772 it = PyObject_GetIter(seq);
773 if (it == NULL)
774 return NULL;
775
776 /* create dropwhileobject structure */
777 lz = (dropwhileobject *)type->tp_alloc(type, 0);
778 if (lz == NULL) {
779 Py_DECREF(it);
780 return NULL;
781 }
782 Py_INCREF(func);
783 lz->func = func;
784 lz->it = it;
785 lz->start = 0;
786
787 return (PyObject *)lz;
788}
789
790static void
791dropwhile_dealloc(dropwhileobject *lz)
792{
793 PyObject_GC_UnTrack(lz);
794 Py_XDECREF(lz->func);
795 Py_XDECREF(lz->it);
796 lz->ob_type->tp_free(lz);
797}
798
799static int
800dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
801{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000802 Py_VISIT(lz->it);
803 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000804 return 0;
805}
806
807static PyObject *
808dropwhile_next(dropwhileobject *lz)
809{
810 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000811 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000812 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000813 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000814
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000815 assert(PyIter_Check(it));
816 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000817 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000818 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000819 if (item == NULL)
820 return NULL;
821 if (lz->start == 1)
822 return item;
823
824 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
825 if (good == NULL) {
826 Py_DECREF(item);
827 return NULL;
828 }
829 ok = PyObject_IsTrue(good);
830 Py_DECREF(good);
831 if (!ok) {
832 lz->start = 1;
833 return item;
834 }
835 Py_DECREF(item);
836 }
837}
838
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000839PyDoc_STRVAR(dropwhile_doc,
840"dropwhile(predicate, iterable) --> dropwhile object\n\
841\n\
842Drop items from the iterable while predicate(item) is true.\n\
843Afterwards, return every element until the iterable is exhausted.");
844
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000845static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000846 PyObject_HEAD_INIT(NULL)
847 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000848 "itertools.dropwhile", /* tp_name */
849 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000850 0, /* tp_itemsize */
851 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000852 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000853 0, /* tp_print */
854 0, /* tp_getattr */
855 0, /* tp_setattr */
856 0, /* tp_compare */
857 0, /* tp_repr */
858 0, /* tp_as_number */
859 0, /* tp_as_sequence */
860 0, /* tp_as_mapping */
861 0, /* tp_hash */
862 0, /* tp_call */
863 0, /* tp_str */
864 PyObject_GenericGetAttr, /* tp_getattro */
865 0, /* tp_setattro */
866 0, /* tp_as_buffer */
867 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
868 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000869 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000870 (traverseproc)dropwhile_traverse, /* tp_traverse */
871 0, /* tp_clear */
872 0, /* tp_richcompare */
873 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000874 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000875 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876 0, /* tp_methods */
877 0, /* tp_members */
878 0, /* tp_getset */
879 0, /* tp_base */
880 0, /* tp_dict */
881 0, /* tp_descr_get */
882 0, /* tp_descr_set */
883 0, /* tp_dictoffset */
884 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000885 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000886 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000887 PyObject_GC_Del, /* tp_free */
888};
889
890
891/* takewhile object **********************************************************/
892
893typedef struct {
894 PyObject_HEAD
895 PyObject *func;
896 PyObject *it;
897 long stop;
898} takewhileobject;
899
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000900static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000901
902static PyObject *
903takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
904{
905 PyObject *func, *seq;
906 PyObject *it;
907 takewhileobject *lz;
908
909 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
910 return NULL;
911
912 /* Get iterator. */
913 it = PyObject_GetIter(seq);
914 if (it == NULL)
915 return NULL;
916
917 /* create takewhileobject structure */
918 lz = (takewhileobject *)type->tp_alloc(type, 0);
919 if (lz == NULL) {
920 Py_DECREF(it);
921 return NULL;
922 }
923 Py_INCREF(func);
924 lz->func = func;
925 lz->it = it;
926 lz->stop = 0;
927
928 return (PyObject *)lz;
929}
930
931static void
932takewhile_dealloc(takewhileobject *lz)
933{
934 PyObject_GC_UnTrack(lz);
935 Py_XDECREF(lz->func);
936 Py_XDECREF(lz->it);
937 lz->ob_type->tp_free(lz);
938}
939
940static int
941takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
942{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000943 Py_VISIT(lz->it);
944 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000945 return 0;
946}
947
948static PyObject *
949takewhile_next(takewhileobject *lz)
950{
951 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000952 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000953 long ok;
954
955 if (lz->stop == 1)
956 return NULL;
957
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000958 assert(PyIter_Check(it));
959 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000960 if (item == NULL)
961 return NULL;
962
963 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
964 if (good == NULL) {
965 Py_DECREF(item);
966 return NULL;
967 }
968 ok = PyObject_IsTrue(good);
969 Py_DECREF(good);
970 if (ok)
971 return item;
972 Py_DECREF(item);
973 lz->stop = 1;
974 return NULL;
975}
976
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000977PyDoc_STRVAR(takewhile_doc,
978"takewhile(predicate, iterable) --> takewhile object\n\
979\n\
980Return successive entries from an iterable as long as the \n\
981predicate evaluates to true for each entry.");
982
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000983static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000984 PyObject_HEAD_INIT(NULL)
985 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000986 "itertools.takewhile", /* tp_name */
987 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000988 0, /* tp_itemsize */
989 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000990 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000991 0, /* tp_print */
992 0, /* tp_getattr */
993 0, /* tp_setattr */
994 0, /* tp_compare */
995 0, /* tp_repr */
996 0, /* tp_as_number */
997 0, /* tp_as_sequence */
998 0, /* tp_as_mapping */
999 0, /* tp_hash */
1000 0, /* tp_call */
1001 0, /* tp_str */
1002 PyObject_GenericGetAttr, /* tp_getattro */
1003 0, /* tp_setattro */
1004 0, /* tp_as_buffer */
1005 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1006 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001007 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001008 (traverseproc)takewhile_traverse, /* tp_traverse */
1009 0, /* tp_clear */
1010 0, /* tp_richcompare */
1011 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001012 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001013 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014 0, /* tp_methods */
1015 0, /* tp_members */
1016 0, /* tp_getset */
1017 0, /* tp_base */
1018 0, /* tp_dict */
1019 0, /* tp_descr_get */
1020 0, /* tp_descr_set */
1021 0, /* tp_dictoffset */
1022 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001023 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001024 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001025 PyObject_GC_Del, /* tp_free */
1026};
1027
1028
1029/* islice object ************************************************************/
1030
1031typedef struct {
1032 PyObject_HEAD
1033 PyObject *it;
1034 long next;
1035 long stop;
1036 long step;
1037 long cnt;
1038} isliceobject;
1039
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001040static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001041
1042static PyObject *
1043islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1044{
1045 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001046 long start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001047 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001048 int numargs;
1049 isliceobject *lz;
1050
Raymond Hettingerb2594052004-12-05 09:25:51 +00001051 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052 return NULL;
1053
Raymond Hettingerb2594052004-12-05 09:25:51 +00001054 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001055 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001056 if (a1 != Py_None) {
1057 stop = PyInt_AsLong(a1);
1058 if (stop == -1) {
1059 if (PyErr_Occurred())
1060 PyErr_Clear();
1061 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001062 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001063 return NULL;
1064 }
1065 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001066 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001067 if (a1 != Py_None)
1068 start = PyInt_AsLong(a1);
1069 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001070 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001071 if (a2 != Py_None) {
1072 stop = PyInt_AsLong(a2);
1073 if (stop == -1) {
1074 if (PyErr_Occurred())
1075 PyErr_Clear();
1076 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001077 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001078 return NULL;
1079 }
1080 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001082 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001083 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001084 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001085 return NULL;
1086 }
1087
Raymond Hettingerb2594052004-12-05 09:25:51 +00001088 if (a3 != NULL) {
1089 if (a3 != Py_None)
1090 step = PyInt_AsLong(a3);
1091 if (step == -1 && PyErr_Occurred())
1092 PyErr_Clear();
1093 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001094 if (step<1) {
1095 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001096 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001097 return NULL;
1098 }
1099
1100 /* Get iterator. */
1101 it = PyObject_GetIter(seq);
1102 if (it == NULL)
1103 return NULL;
1104
1105 /* create isliceobject structure */
1106 lz = (isliceobject *)type->tp_alloc(type, 0);
1107 if (lz == NULL) {
1108 Py_DECREF(it);
1109 return NULL;
1110 }
1111 lz->it = it;
1112 lz->next = start;
1113 lz->stop = stop;
1114 lz->step = step;
1115 lz->cnt = 0L;
1116
1117 return (PyObject *)lz;
1118}
1119
1120static void
1121islice_dealloc(isliceobject *lz)
1122{
1123 PyObject_GC_UnTrack(lz);
1124 Py_XDECREF(lz->it);
1125 lz->ob_type->tp_free(lz);
1126}
1127
1128static int
1129islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1130{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001131 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001132 return 0;
1133}
1134
1135static PyObject *
1136islice_next(isliceobject *lz)
1137{
1138 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001139 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001140 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001141 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001142
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001143 assert(PyIter_Check(it));
1144 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001145 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001146 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001147 if (item == NULL)
1148 return NULL;
1149 Py_DECREF(item);
1150 lz->cnt++;
1151 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001152 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001154 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001155 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001156 if (item == NULL)
1157 return NULL;
1158 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001159 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001161 if (lz->next < oldnext) /* Check for overflow */
1162 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163 return item;
1164}
1165
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001166PyDoc_STRVAR(islice_doc,
1167"islice(iterable, [start,] stop [, step]) --> islice object\n\
1168\n\
1169Return an iterator whose next() method returns selected values from an\n\
1170iterable. If start is specified, will skip all preceding elements;\n\
1171otherwise, start defaults to zero. Step defaults to one. If\n\
1172specified as another value, step determines how many values are \n\
1173skipped between successive calls. Works like a slice() on a list\n\
1174but returns an iterator.");
1175
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001176static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177 PyObject_HEAD_INIT(NULL)
1178 0, /* ob_size */
1179 "itertools.islice", /* tp_name */
1180 sizeof(isliceobject), /* tp_basicsize */
1181 0, /* tp_itemsize */
1182 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001183 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001184 0, /* tp_print */
1185 0, /* tp_getattr */
1186 0, /* tp_setattr */
1187 0, /* tp_compare */
1188 0, /* tp_repr */
1189 0, /* tp_as_number */
1190 0, /* tp_as_sequence */
1191 0, /* tp_as_mapping */
1192 0, /* tp_hash */
1193 0, /* tp_call */
1194 0, /* tp_str */
1195 PyObject_GenericGetAttr, /* tp_getattro */
1196 0, /* tp_setattro */
1197 0, /* tp_as_buffer */
1198 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1199 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001200 islice_doc, /* tp_doc */
1201 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001202 0, /* tp_clear */
1203 0, /* tp_richcompare */
1204 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001205 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001206 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207 0, /* tp_methods */
1208 0, /* tp_members */
1209 0, /* tp_getset */
1210 0, /* tp_base */
1211 0, /* tp_dict */
1212 0, /* tp_descr_get */
1213 0, /* tp_descr_set */
1214 0, /* tp_dictoffset */
1215 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001216 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217 islice_new, /* tp_new */
1218 PyObject_GC_Del, /* tp_free */
1219};
1220
1221
1222/* starmap object ************************************************************/
1223
1224typedef struct {
1225 PyObject_HEAD
1226 PyObject *func;
1227 PyObject *it;
1228} starmapobject;
1229
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001230static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231
1232static PyObject *
1233starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1234{
1235 PyObject *func, *seq;
1236 PyObject *it;
1237 starmapobject *lz;
1238
1239 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1240 return NULL;
1241
1242 /* Get iterator. */
1243 it = PyObject_GetIter(seq);
1244 if (it == NULL)
1245 return NULL;
1246
1247 /* create starmapobject structure */
1248 lz = (starmapobject *)type->tp_alloc(type, 0);
1249 if (lz == NULL) {
1250 Py_DECREF(it);
1251 return NULL;
1252 }
1253 Py_INCREF(func);
1254 lz->func = func;
1255 lz->it = it;
1256
1257 return (PyObject *)lz;
1258}
1259
1260static void
1261starmap_dealloc(starmapobject *lz)
1262{
1263 PyObject_GC_UnTrack(lz);
1264 Py_XDECREF(lz->func);
1265 Py_XDECREF(lz->it);
1266 lz->ob_type->tp_free(lz);
1267}
1268
1269static int
1270starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1271{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001272 Py_VISIT(lz->it);
1273 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001274 return 0;
1275}
1276
1277static PyObject *
1278starmap_next(starmapobject *lz)
1279{
1280 PyObject *args;
1281 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001282 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001283
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001284 assert(PyIter_Check(it));
1285 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001286 if (args == NULL)
1287 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001288 if (!PyTuple_CheckExact(args)) {
1289 Py_DECREF(args);
1290 PyErr_SetString(PyExc_TypeError,
1291 "iterator must return a tuple");
1292 return NULL;
1293 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294 result = PyObject_Call(lz->func, args, NULL);
1295 Py_DECREF(args);
1296 return result;
1297}
1298
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001299PyDoc_STRVAR(starmap_doc,
1300"starmap(function, sequence) --> starmap object\n\
1301\n\
1302Return an iterator whose values are returned from the function evaluated\n\
1303with a argument tuple taken from the given sequence.");
1304
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001305static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306 PyObject_HEAD_INIT(NULL)
1307 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001308 "itertools.starmap", /* tp_name */
1309 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001310 0, /* tp_itemsize */
1311 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001312 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001313 0, /* tp_print */
1314 0, /* tp_getattr */
1315 0, /* tp_setattr */
1316 0, /* tp_compare */
1317 0, /* tp_repr */
1318 0, /* tp_as_number */
1319 0, /* tp_as_sequence */
1320 0, /* tp_as_mapping */
1321 0, /* tp_hash */
1322 0, /* tp_call */
1323 0, /* tp_str */
1324 PyObject_GenericGetAttr, /* tp_getattro */
1325 0, /* tp_setattro */
1326 0, /* tp_as_buffer */
1327 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1328 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001329 starmap_doc, /* tp_doc */
1330 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001331 0, /* tp_clear */
1332 0, /* tp_richcompare */
1333 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001334 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001335 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001336 0, /* tp_methods */
1337 0, /* tp_members */
1338 0, /* tp_getset */
1339 0, /* tp_base */
1340 0, /* tp_dict */
1341 0, /* tp_descr_get */
1342 0, /* tp_descr_set */
1343 0, /* tp_dictoffset */
1344 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001345 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001346 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001347 PyObject_GC_Del, /* tp_free */
1348};
1349
1350
1351/* imap object ************************************************************/
1352
1353typedef struct {
1354 PyObject_HEAD
1355 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001356 PyObject *func;
1357} imapobject;
1358
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001359static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360
1361static PyObject *
1362imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1363{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001364 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001365 imapobject *lz;
1366 int numargs, i;
1367
1368 numargs = PyTuple_Size(args);
1369 if (numargs < 2) {
1370 PyErr_SetString(PyExc_TypeError,
1371 "imap() must have at least two arguments.");
1372 return NULL;
1373 }
1374
1375 iters = PyTuple_New(numargs-1);
1376 if (iters == NULL)
1377 return NULL;
1378
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001379 for (i=1 ; i<numargs ; i++) {
1380 /* Get iterator. */
1381 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1382 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383 Py_DECREF(iters);
1384 return NULL;
1385 }
1386 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387 }
1388
1389 /* create imapobject structure */
1390 lz = (imapobject *)type->tp_alloc(type, 0);
1391 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001392 Py_DECREF(iters);
1393 return NULL;
1394 }
1395 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001396 func = PyTuple_GET_ITEM(args, 0);
1397 Py_INCREF(func);
1398 lz->func = func;
1399
1400 return (PyObject *)lz;
1401}
1402
1403static void
1404imap_dealloc(imapobject *lz)
1405{
1406 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001407 Py_XDECREF(lz->iters);
1408 Py_XDECREF(lz->func);
1409 lz->ob_type->tp_free(lz);
1410}
1411
1412static int
1413imap_traverse(imapobject *lz, visitproc visit, void *arg)
1414{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001415 Py_VISIT(lz->iters);
1416 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001417 return 0;
1418}
1419
Raymond Hettinger2012f172003-02-07 05:32:58 +00001420/*
1421imap() is an iterator version of __builtins__.map() except that it does
1422not have the None fill-in feature. That was intentionally left out for
1423the following reasons:
1424
1425 1) Itertools are designed to be easily combined and chained together.
1426 Having all tools stop with the shortest input is a unifying principle
1427 that makes it easier to combine finite iterators (supplying data) with
1428 infinite iterators like count() and repeat() (for supplying sequential
1429 or constant arguments to a function).
1430
1431 2) In typical use cases for combining itertools, having one finite data
1432 supplier run out before another is likely to be an error condition which
1433 should not pass silently by automatically supplying None.
1434
1435 3) The use cases for automatic None fill-in are rare -- not many functions
1436 do something useful when a parameter suddenly switches type and becomes
1437 None.
1438
1439 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001440 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001441
1442 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1443*/
1444
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001445static PyObject *
1446imap_next(imapobject *lz)
1447{
1448 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001449 PyObject *argtuple;
1450 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001451 int numargs, i;
1452
1453 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001454 argtuple = PyTuple_New(numargs);
1455 if (argtuple == NULL)
1456 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001458 for (i=0 ; i<numargs ; i++) {
1459 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1460 if (val == NULL) {
1461 Py_DECREF(argtuple);
1462 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001464 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001466 if (lz->func == Py_None)
1467 return argtuple;
1468 result = PyObject_Call(lz->func, argtuple, NULL);
1469 Py_DECREF(argtuple);
1470 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471}
1472
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473PyDoc_STRVAR(imap_doc,
1474"imap(func, *iterables) --> imap object\n\
1475\n\
1476Make an iterator that computes the function using arguments from\n\
1477each of the iterables. Like map() except that it returns\n\
1478an iterator instead of a list and that it stops when the shortest\n\
1479iterable is exhausted instead of filling in None for shorter\n\
1480iterables.");
1481
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001482static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483 PyObject_HEAD_INIT(NULL)
1484 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001485 "itertools.imap", /* tp_name */
1486 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487 0, /* tp_itemsize */
1488 /* methods */
1489 (destructor)imap_dealloc, /* tp_dealloc */
1490 0, /* tp_print */
1491 0, /* tp_getattr */
1492 0, /* tp_setattr */
1493 0, /* tp_compare */
1494 0, /* tp_repr */
1495 0, /* tp_as_number */
1496 0, /* tp_as_sequence */
1497 0, /* tp_as_mapping */
1498 0, /* tp_hash */
1499 0, /* tp_call */
1500 0, /* tp_str */
1501 PyObject_GenericGetAttr, /* tp_getattro */
1502 0, /* tp_setattro */
1503 0, /* tp_as_buffer */
1504 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1505 Py_TPFLAGS_BASETYPE, /* tp_flags */
1506 imap_doc, /* tp_doc */
1507 (traverseproc)imap_traverse, /* tp_traverse */
1508 0, /* tp_clear */
1509 0, /* tp_richcompare */
1510 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001511 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001512 (iternextfunc)imap_next, /* tp_iternext */
1513 0, /* tp_methods */
1514 0, /* tp_members */
1515 0, /* tp_getset */
1516 0, /* tp_base */
1517 0, /* tp_dict */
1518 0, /* tp_descr_get */
1519 0, /* tp_descr_set */
1520 0, /* tp_dictoffset */
1521 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001522 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001523 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524 PyObject_GC_Del, /* tp_free */
1525};
1526
1527
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001528/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001529
1530typedef struct {
1531 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001532 long tuplesize;
1533 long iternum; /* which iterator is active */
1534 PyObject *ittuple; /* tuple of iterators */
1535} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001537static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538
1539static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001540chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001542 chainobject *lz;
1543 int tuplesize = PySequence_Length(args);
1544 int i;
1545 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001546
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001547 /* obtain iterators */
1548 assert(PyTuple_Check(args));
1549 ittuple = PyTuple_New(tuplesize);
1550 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001552 for (i=0; i < tuplesize; ++i) {
1553 PyObject *item = PyTuple_GET_ITEM(args, i);
1554 PyObject *it = PyObject_GetIter(item);
1555 if (it == NULL) {
1556 if (PyErr_ExceptionMatches(PyExc_TypeError))
1557 PyErr_Format(PyExc_TypeError,
1558 "chain argument #%d must support iteration",
1559 i+1);
1560 Py_DECREF(ittuple);
1561 return NULL;
1562 }
1563 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564 }
1565
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001566 /* create chainobject structure */
1567 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001568 if (lz == NULL) {
1569 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001570 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001571 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001572
1573 lz->ittuple = ittuple;
1574 lz->iternum = 0;
1575 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001576
1577 return (PyObject *)lz;
1578}
1579
1580static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001581chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001582{
1583 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001584 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001585 lz->ob_type->tp_free(lz);
1586}
1587
Raymond Hettinger2012f172003-02-07 05:32:58 +00001588static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001589chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001590{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001591 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001592 return 0;
1593}
1594
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001596chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001597{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001598 PyObject *it;
1599 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001600
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001601 while (lz->iternum < lz->tuplesize) {
1602 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1603 item = PyIter_Next(it);
1604 if (item != NULL)
1605 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001606 if (PyErr_Occurred()) {
1607 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1608 PyErr_Clear();
1609 else
1610 return NULL;
1611 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001612 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001613 }
1614 return NULL;
1615}
1616
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001617PyDoc_STRVAR(chain_doc,
1618"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001620Return a chain object whose .next() method returns elements from the\n\
1621first iterable until it is exhausted, then elements from the next\n\
1622iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001623
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001624static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001625 PyObject_HEAD_INIT(NULL)
1626 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001627 "itertools.chain", /* tp_name */
1628 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001629 0, /* tp_itemsize */
1630 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001631 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001632 0, /* tp_print */
1633 0, /* tp_getattr */
1634 0, /* tp_setattr */
1635 0, /* tp_compare */
1636 0, /* tp_repr */
1637 0, /* tp_as_number */
1638 0, /* tp_as_sequence */
1639 0, /* tp_as_mapping */
1640 0, /* tp_hash */
1641 0, /* tp_call */
1642 0, /* tp_str */
1643 PyObject_GenericGetAttr, /* tp_getattro */
1644 0, /* tp_setattro */
1645 0, /* tp_as_buffer */
1646 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1647 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001648 chain_doc, /* tp_doc */
1649 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650 0, /* tp_clear */
1651 0, /* tp_richcompare */
1652 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001653 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001654 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001655 0, /* tp_methods */
1656 0, /* tp_members */
1657 0, /* tp_getset */
1658 0, /* tp_base */
1659 0, /* tp_dict */
1660 0, /* tp_descr_get */
1661 0, /* tp_descr_set */
1662 0, /* tp_dictoffset */
1663 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001664 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001665 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666 PyObject_GC_Del, /* tp_free */
1667};
1668
1669
1670/* ifilter object ************************************************************/
1671
1672typedef struct {
1673 PyObject_HEAD
1674 PyObject *func;
1675 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676} ifilterobject;
1677
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001678static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001679
1680static PyObject *
1681ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1682{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001683 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684 PyObject *it;
1685 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001686
Raymond Hettinger60eca932003-02-09 06:40:58 +00001687 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001688 return NULL;
1689
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690 /* Get iterator. */
1691 it = PyObject_GetIter(seq);
1692 if (it == NULL)
1693 return NULL;
1694
1695 /* create ifilterobject structure */
1696 lz = (ifilterobject *)type->tp_alloc(type, 0);
1697 if (lz == NULL) {
1698 Py_DECREF(it);
1699 return NULL;
1700 }
1701 Py_INCREF(func);
1702 lz->func = func;
1703 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001704
1705 return (PyObject *)lz;
1706}
1707
1708static void
1709ifilter_dealloc(ifilterobject *lz)
1710{
1711 PyObject_GC_UnTrack(lz);
1712 Py_XDECREF(lz->func);
1713 Py_XDECREF(lz->it);
1714 lz->ob_type->tp_free(lz);
1715}
1716
1717static int
1718ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1719{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001720 Py_VISIT(lz->it);
1721 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001722 return 0;
1723}
1724
1725static PyObject *
1726ifilter_next(ifilterobject *lz)
1727{
1728 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001729 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001730 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001731 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001732
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001733 assert(PyIter_Check(it));
1734 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001736 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737 if (item == NULL)
1738 return NULL;
1739
1740 if (lz->func == Py_None) {
1741 ok = PyObject_IsTrue(item);
1742 } else {
1743 PyObject *good;
1744 good = PyObject_CallFunctionObjArgs(lz->func,
1745 item, NULL);
1746 if (good == NULL) {
1747 Py_DECREF(item);
1748 return NULL;
1749 }
1750 ok = PyObject_IsTrue(good);
1751 Py_DECREF(good);
1752 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001753 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001754 return item;
1755 Py_DECREF(item);
1756 }
1757}
1758
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001760"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001761\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001762Return those items of sequence for which function(item) is true.\n\
1763If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001765static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001766 PyObject_HEAD_INIT(NULL)
1767 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001768 "itertools.ifilter", /* tp_name */
1769 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001770 0, /* tp_itemsize */
1771 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001772 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001773 0, /* tp_print */
1774 0, /* tp_getattr */
1775 0, /* tp_setattr */
1776 0, /* tp_compare */
1777 0, /* tp_repr */
1778 0, /* tp_as_number */
1779 0, /* tp_as_sequence */
1780 0, /* tp_as_mapping */
1781 0, /* tp_hash */
1782 0, /* tp_call */
1783 0, /* tp_str */
1784 PyObject_GenericGetAttr, /* tp_getattro */
1785 0, /* tp_setattro */
1786 0, /* tp_as_buffer */
1787 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1788 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001789 ifilter_doc, /* tp_doc */
1790 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001791 0, /* tp_clear */
1792 0, /* tp_richcompare */
1793 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001794 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001795 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001796 0, /* tp_methods */
1797 0, /* tp_members */
1798 0, /* tp_getset */
1799 0, /* tp_base */
1800 0, /* tp_dict */
1801 0, /* tp_descr_get */
1802 0, /* tp_descr_set */
1803 0, /* tp_dictoffset */
1804 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001805 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001806 ifilter_new, /* tp_new */
1807 PyObject_GC_Del, /* tp_free */
1808};
1809
1810
1811/* ifilterfalse object ************************************************************/
1812
1813typedef struct {
1814 PyObject_HEAD
1815 PyObject *func;
1816 PyObject *it;
1817} ifilterfalseobject;
1818
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001819static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001820
1821static PyObject *
1822ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1823{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001824 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001825 PyObject *it;
1826 ifilterfalseobject *lz;
1827
1828 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1829 return NULL;
1830
1831 /* Get iterator. */
1832 it = PyObject_GetIter(seq);
1833 if (it == NULL)
1834 return NULL;
1835
1836 /* create ifilterfalseobject structure */
1837 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1838 if (lz == NULL) {
1839 Py_DECREF(it);
1840 return NULL;
1841 }
1842 Py_INCREF(func);
1843 lz->func = func;
1844 lz->it = it;
1845
1846 return (PyObject *)lz;
1847}
1848
1849static void
1850ifilterfalse_dealloc(ifilterfalseobject *lz)
1851{
1852 PyObject_GC_UnTrack(lz);
1853 Py_XDECREF(lz->func);
1854 Py_XDECREF(lz->it);
1855 lz->ob_type->tp_free(lz);
1856}
1857
1858static int
1859ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1860{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001861 Py_VISIT(lz->it);
1862 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001863 return 0;
1864}
1865
1866static PyObject *
1867ifilterfalse_next(ifilterfalseobject *lz)
1868{
1869 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001870 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001871 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001872 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001873
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001874 assert(PyIter_Check(it));
1875 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001876 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001877 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001878 if (item == NULL)
1879 return NULL;
1880
1881 if (lz->func == Py_None) {
1882 ok = PyObject_IsTrue(item);
1883 } else {
1884 PyObject *good;
1885 good = PyObject_CallFunctionObjArgs(lz->func,
1886 item, NULL);
1887 if (good == NULL) {
1888 Py_DECREF(item);
1889 return NULL;
1890 }
1891 ok = PyObject_IsTrue(good);
1892 Py_DECREF(good);
1893 }
1894 if (!ok)
1895 return item;
1896 Py_DECREF(item);
1897 }
1898}
1899
Raymond Hettinger60eca932003-02-09 06:40:58 +00001900PyDoc_STRVAR(ifilterfalse_doc,
1901"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1902\n\
1903Return those items of sequence for which function(item) is false.\n\
1904If function is None, return the items that are false.");
1905
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001906static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001907 PyObject_HEAD_INIT(NULL)
1908 0, /* ob_size */
1909 "itertools.ifilterfalse", /* tp_name */
1910 sizeof(ifilterfalseobject), /* tp_basicsize */
1911 0, /* tp_itemsize */
1912 /* methods */
1913 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1914 0, /* tp_print */
1915 0, /* tp_getattr */
1916 0, /* tp_setattr */
1917 0, /* tp_compare */
1918 0, /* tp_repr */
1919 0, /* tp_as_number */
1920 0, /* tp_as_sequence */
1921 0, /* tp_as_mapping */
1922 0, /* tp_hash */
1923 0, /* tp_call */
1924 0, /* tp_str */
1925 PyObject_GenericGetAttr, /* tp_getattro */
1926 0, /* tp_setattro */
1927 0, /* tp_as_buffer */
1928 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1929 Py_TPFLAGS_BASETYPE, /* tp_flags */
1930 ifilterfalse_doc, /* tp_doc */
1931 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1932 0, /* tp_clear */
1933 0, /* tp_richcompare */
1934 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001935 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001936 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1937 0, /* tp_methods */
1938 0, /* tp_members */
1939 0, /* tp_getset */
1940 0, /* tp_base */
1941 0, /* tp_dict */
1942 0, /* tp_descr_get */
1943 0, /* tp_descr_set */
1944 0, /* tp_dictoffset */
1945 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001946 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001947 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001948 PyObject_GC_Del, /* tp_free */
1949};
1950
1951
1952/* count object ************************************************************/
1953
1954typedef struct {
1955 PyObject_HEAD
1956 long cnt;
1957} countobject;
1958
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001959static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001960
1961static PyObject *
1962count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1963{
1964 countobject *lz;
1965 long cnt = 0;
1966
1967 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1968 return NULL;
1969
1970 /* create countobject structure */
1971 lz = (countobject *)PyObject_New(countobject, &count_type);
1972 if (lz == NULL)
1973 return NULL;
1974 lz->cnt = cnt;
1975
1976 return (PyObject *)lz;
1977}
1978
1979static PyObject *
1980count_next(countobject *lz)
1981{
1982 return PyInt_FromLong(lz->cnt++);
1983}
1984
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001985static PyObject *
1986count_repr(countobject *lz)
1987{
Brett Cannon00468392004-04-13 02:43:53 +00001988 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001989}
1990
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001991PyDoc_STRVAR(count_doc,
1992"count([firstval]) --> count object\n\
1993\n\
1994Return a count object whose .next() method returns consecutive\n\
1995integers starting from zero or, if specified, from firstval.");
1996
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001997static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001998 PyObject_HEAD_INIT(NULL)
1999 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002000 "itertools.count", /* tp_name */
2001 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002002 0, /* tp_itemsize */
2003 /* methods */
2004 (destructor)PyObject_Del, /* tp_dealloc */
2005 0, /* tp_print */
2006 0, /* tp_getattr */
2007 0, /* tp_setattr */
2008 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002009 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002010 0, /* tp_as_number */
2011 0, /* tp_as_sequence */
2012 0, /* tp_as_mapping */
2013 0, /* tp_hash */
2014 0, /* tp_call */
2015 0, /* tp_str */
2016 PyObject_GenericGetAttr, /* tp_getattro */
2017 0, /* tp_setattro */
2018 0, /* tp_as_buffer */
2019 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002020 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002021 0, /* tp_traverse */
2022 0, /* tp_clear */
2023 0, /* tp_richcompare */
2024 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002025 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002026 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002027 0, /* tp_methods */
2028 0, /* tp_members */
2029 0, /* tp_getset */
2030 0, /* tp_base */
2031 0, /* tp_dict */
2032 0, /* tp_descr_get */
2033 0, /* tp_descr_set */
2034 0, /* tp_dictoffset */
2035 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002036 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002037 count_new, /* tp_new */
2038};
2039
2040
2041/* izip object ************************************************************/
2042
2043#include "Python.h"
2044
2045typedef struct {
2046 PyObject_HEAD
2047 long tuplesize;
2048 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002049 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002050} izipobject;
2051
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002052static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002053
2054static PyObject *
2055izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2056{
2057 izipobject *lz;
2058 int i;
2059 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002060 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002061 int tuplesize = PySequence_Length(args);
2062
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002063 /* args must be a tuple */
2064 assert(PyTuple_Check(args));
2065
2066 /* obtain iterators */
2067 ittuple = PyTuple_New(tuplesize);
2068 if(ittuple == NULL)
2069 return NULL;
2070 for (i=0; i < tuplesize; ++i) {
2071 PyObject *item = PyTuple_GET_ITEM(args, i);
2072 PyObject *it = PyObject_GetIter(item);
2073 if (it == NULL) {
2074 if (PyErr_ExceptionMatches(PyExc_TypeError))
2075 PyErr_Format(PyExc_TypeError,
2076 "izip argument #%d must support iteration",
2077 i+1);
2078 Py_DECREF(ittuple);
2079 return NULL;
2080 }
2081 PyTuple_SET_ITEM(ittuple, i, it);
2082 }
2083
Raymond Hettinger2012f172003-02-07 05:32:58 +00002084 /* create a result holder */
2085 result = PyTuple_New(tuplesize);
2086 if (result == NULL) {
2087 Py_DECREF(ittuple);
2088 return NULL;
2089 }
2090 for (i=0 ; i < tuplesize ; i++) {
2091 Py_INCREF(Py_None);
2092 PyTuple_SET_ITEM(result, i, Py_None);
2093 }
2094
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002095 /* create izipobject structure */
2096 lz = (izipobject *)type->tp_alloc(type, 0);
2097 if (lz == NULL) {
2098 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002099 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002100 return NULL;
2101 }
2102 lz->ittuple = ittuple;
2103 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002104 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002105
2106 return (PyObject *)lz;
2107}
2108
2109static void
2110izip_dealloc(izipobject *lz)
2111{
2112 PyObject_GC_UnTrack(lz);
2113 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002114 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002115 lz->ob_type->tp_free(lz);
2116}
2117
2118static int
2119izip_traverse(izipobject *lz, visitproc visit, void *arg)
2120{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002121 Py_VISIT(lz->ittuple);
2122 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002123 return 0;
2124}
2125
2126static PyObject *
2127izip_next(izipobject *lz)
2128{
2129 int i;
2130 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002131 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002132 PyObject *it;
2133 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002134 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002135
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002136 if (tuplesize == 0)
2137 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002138 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002139 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002140 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002141 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002142 assert(PyIter_Check(it));
2143 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002144 if (item == NULL) {
2145 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002146 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002147 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002148 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002149 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002150 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002151 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002152 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002153 result = PyTuple_New(tuplesize);
2154 if (result == NULL)
2155 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002156 for (i=0 ; i < tuplesize ; i++) {
2157 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002158 assert(PyIter_Check(it));
2159 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002160 if (item == NULL) {
2161 Py_DECREF(result);
2162 return NULL;
2163 }
2164 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002165 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002166 }
2167 return result;
2168}
2169
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002170PyDoc_STRVAR(izip_doc,
2171"izip(iter1 [,iter2 [...]]) --> izip object\n\
2172\n\
2173Return a izip object whose .next() method returns a tuple where\n\
2174the i-th element comes from the i-th iterable argument. The .next()\n\
2175method continues until the shortest iterable in the argument sequence\n\
2176is exhausted and then it raises StopIteration. Works like the zip()\n\
2177function but consumes less memory by returning an iterator instead of\n\
2178a list.");
2179
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002180static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002181 PyObject_HEAD_INIT(NULL)
2182 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002183 "itertools.izip", /* tp_name */
2184 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002185 0, /* tp_itemsize */
2186 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002187 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002188 0, /* tp_print */
2189 0, /* tp_getattr */
2190 0, /* tp_setattr */
2191 0, /* tp_compare */
2192 0, /* tp_repr */
2193 0, /* tp_as_number */
2194 0, /* tp_as_sequence */
2195 0, /* tp_as_mapping */
2196 0, /* tp_hash */
2197 0, /* tp_call */
2198 0, /* tp_str */
2199 PyObject_GenericGetAttr, /* tp_getattro */
2200 0, /* tp_setattro */
2201 0, /* tp_as_buffer */
2202 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2203 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002204 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002205 (traverseproc)izip_traverse, /* tp_traverse */
2206 0, /* tp_clear */
2207 0, /* tp_richcompare */
2208 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002209 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002210 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002211 0, /* tp_methods */
2212 0, /* tp_members */
2213 0, /* tp_getset */
2214 0, /* tp_base */
2215 0, /* tp_dict */
2216 0, /* tp_descr_get */
2217 0, /* tp_descr_set */
2218 0, /* tp_dictoffset */
2219 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002220 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002221 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002222 PyObject_GC_Del, /* tp_free */
2223};
2224
2225
2226/* repeat object ************************************************************/
2227
2228typedef struct {
2229 PyObject_HEAD
2230 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002231 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002232} repeatobject;
2233
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002234static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002235
2236static PyObject *
2237repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2238{
2239 repeatobject *ro;
2240 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002241 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002242
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002243 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002244 return NULL;
2245
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002246 if (PyTuple_Size(args) == 2 && cnt < 0)
2247 cnt = 0;
2248
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002249 ro = (repeatobject *)type->tp_alloc(type, 0);
2250 if (ro == NULL)
2251 return NULL;
2252 Py_INCREF(element);
2253 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002254 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002255 return (PyObject *)ro;
2256}
2257
2258static void
2259repeat_dealloc(repeatobject *ro)
2260{
2261 PyObject_GC_UnTrack(ro);
2262 Py_XDECREF(ro->element);
2263 ro->ob_type->tp_free(ro);
2264}
2265
2266static int
2267repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2268{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002269 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002270 return 0;
2271}
2272
2273static PyObject *
2274repeat_next(repeatobject *ro)
2275{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002276 if (ro->cnt == 0)
2277 return NULL;
2278 if (ro->cnt > 0)
2279 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002280 Py_INCREF(ro->element);
2281 return ro->element;
2282}
2283
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002284static PyObject *
2285repeat_repr(repeatobject *ro)
2286{
2287 PyObject *result, *objrepr;
2288
2289 objrepr = PyObject_Repr(ro->element);
2290 if (objrepr == NULL)
2291 return NULL;
2292
2293 if (ro->cnt == -1)
2294 result = PyString_FromFormat("repeat(%s)",
2295 PyString_AS_STRING(objrepr));
2296 else
Brett Cannon00468392004-04-13 02:43:53 +00002297 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002298 PyString_AS_STRING(objrepr), ro->cnt);
2299 Py_DECREF(objrepr);
2300 return result;
2301}
2302
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002303static int
2304repeat_len(repeatobject *ro)
2305{
2306 if (ro->cnt == -1)
2307 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2308 return (int)(ro->cnt);
2309}
2310
2311static PySequenceMethods repeat_as_sequence = {
2312 (inquiry)repeat_len, /* sq_length */
2313 0, /* sq_concat */
2314};
2315
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002316PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002317"repeat(element [,times]) -> create an iterator which returns the element\n\
2318for the specified number of times. If not specified, returns the element\n\
2319endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002320
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002321static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002322 PyObject_HEAD_INIT(NULL)
2323 0, /* ob_size */
2324 "itertools.repeat", /* tp_name */
2325 sizeof(repeatobject), /* tp_basicsize */
2326 0, /* tp_itemsize */
2327 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002328 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002329 0, /* tp_print */
2330 0, /* tp_getattr */
2331 0, /* tp_setattr */
2332 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002333 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002334 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002335 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002336 0, /* tp_as_mapping */
2337 0, /* tp_hash */
2338 0, /* tp_call */
2339 0, /* tp_str */
2340 PyObject_GenericGetAttr, /* tp_getattro */
2341 0, /* tp_setattro */
2342 0, /* tp_as_buffer */
2343 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2344 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002345 repeat_doc, /* tp_doc */
2346 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002347 0, /* tp_clear */
2348 0, /* tp_richcompare */
2349 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002350 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002351 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002352 0, /* tp_methods */
2353 0, /* tp_members */
2354 0, /* tp_getset */
2355 0, /* tp_base */
2356 0, /* tp_dict */
2357 0, /* tp_descr_get */
2358 0, /* tp_descr_set */
2359 0, /* tp_dictoffset */
2360 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002361 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002362 repeat_new, /* tp_new */
2363 PyObject_GC_Del, /* tp_free */
2364};
2365
2366
2367/* module level code ********************************************************/
2368
2369PyDoc_STRVAR(module_doc,
2370"Functional tools for creating and using iterators.\n\
2371\n\
2372Infinite iterators:\n\
2373count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002374cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002375repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002376\n\
2377Iterators terminating on the shortest input sequence:\n\
2378izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002379ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2380ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002381islice(seq, [start,] stop [, step]) --> elements from\n\
2382 seq[start:stop:step]\n\
2383imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2384starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002385tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002386chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002387takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2388dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002389groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002390");
2391
2392
Raymond Hettingerad983e72003-11-12 14:32:26 +00002393static PyMethodDef module_methods[] = {
2394 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2395 {NULL, NULL} /* sentinel */
2396};
2397
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002398PyMODINIT_FUNC
2399inititertools(void)
2400{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002401 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002402 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002403 char *name;
2404 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002405 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002406 &dropwhile_type,
2407 &takewhile_type,
2408 &islice_type,
2409 &starmap_type,
2410 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002411 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002412 &ifilter_type,
2413 &ifilterfalse_type,
2414 &count_type,
2415 &izip_type,
2416 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002417 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002418 NULL
2419 };
2420
Skip Montanarof3938fd2004-02-10 20:27:40 +00002421 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002422 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002423
Raymond Hettinger60eca932003-02-09 06:40:58 +00002424 for (i=0 ; typelist[i] != NULL ; i++) {
2425 if (PyType_Ready(typelist[i]) < 0)
2426 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002427 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002428 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002429 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002430 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002431 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002432
2433 if (PyType_Ready(&teedataobject_type) < 0)
2434 return;
2435 if (PyType_Ready(&tee_type) < 0)
2436 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002437 if (PyType_Ready(&_grouper_type) < 0)
2438 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002439}