blob: bf148ac2441bfffdd07d752d33b1ddb0dab90d64 [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;
1047 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001048 int numargs;
1049 isliceobject *lz;
1050
1051 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +00001052 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053 return NULL;
1054
1055 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 Hettinger0e4f7642003-10-28 07:32:28 +00001062 "Stop argument 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 Hettinger14ef54c2003-05-02 19:04:37 +00001067 start = PyInt_AsLong(a1);
1068 if (start == -1 && PyErr_Occurred()) {
1069 PyErr_Clear();
1070 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001071 "Start argument must be a non-negative integer.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001072 return NULL;
1073 }
1074 if (a2 != Py_None) {
1075 stop = PyInt_AsLong(a2);
1076 if (stop == -1) {
1077 if (PyErr_Occurred())
1078 PyErr_Clear();
1079 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001080 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001081 return NULL;
1082 }
1083 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001084 }
1085
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001086 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001087 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001088 "Indices for islice() must be non-negative integers.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001089 return NULL;
1090 }
1091
1092 if (step<1) {
1093 PyErr_SetString(PyExc_ValueError,
1094 "Step must be one or larger for islice().");
1095 return NULL;
1096 }
1097
1098 /* Get iterator. */
1099 it = PyObject_GetIter(seq);
1100 if (it == NULL)
1101 return NULL;
1102
1103 /* create isliceobject structure */
1104 lz = (isliceobject *)type->tp_alloc(type, 0);
1105 if (lz == NULL) {
1106 Py_DECREF(it);
1107 return NULL;
1108 }
1109 lz->it = it;
1110 lz->next = start;
1111 lz->stop = stop;
1112 lz->step = step;
1113 lz->cnt = 0L;
1114
1115 return (PyObject *)lz;
1116}
1117
1118static void
1119islice_dealloc(isliceobject *lz)
1120{
1121 PyObject_GC_UnTrack(lz);
1122 Py_XDECREF(lz->it);
1123 lz->ob_type->tp_free(lz);
1124}
1125
1126static int
1127islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1128{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001129 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001130 return 0;
1131}
1132
1133static PyObject *
1134islice_next(isliceobject *lz)
1135{
1136 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001137 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001138 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001139 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001140
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001141 assert(PyIter_Check(it));
1142 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001143 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001144 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001145 if (item == NULL)
1146 return NULL;
1147 Py_DECREF(item);
1148 lz->cnt++;
1149 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001150 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001151 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001152 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001153 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154 if (item == NULL)
1155 return NULL;
1156 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001157 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001158 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001159 if (lz->next < oldnext) /* Check for overflow */
1160 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 return item;
1162}
1163
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001164PyDoc_STRVAR(islice_doc,
1165"islice(iterable, [start,] stop [, step]) --> islice object\n\
1166\n\
1167Return an iterator whose next() method returns selected values from an\n\
1168iterable. If start is specified, will skip all preceding elements;\n\
1169otherwise, start defaults to zero. Step defaults to one. If\n\
1170specified as another value, step determines how many values are \n\
1171skipped between successive calls. Works like a slice() on a list\n\
1172but returns an iterator.");
1173
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001174static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001175 PyObject_HEAD_INIT(NULL)
1176 0, /* ob_size */
1177 "itertools.islice", /* tp_name */
1178 sizeof(isliceobject), /* tp_basicsize */
1179 0, /* tp_itemsize */
1180 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001181 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001182 0, /* tp_print */
1183 0, /* tp_getattr */
1184 0, /* tp_setattr */
1185 0, /* tp_compare */
1186 0, /* tp_repr */
1187 0, /* tp_as_number */
1188 0, /* tp_as_sequence */
1189 0, /* tp_as_mapping */
1190 0, /* tp_hash */
1191 0, /* tp_call */
1192 0, /* tp_str */
1193 PyObject_GenericGetAttr, /* tp_getattro */
1194 0, /* tp_setattro */
1195 0, /* tp_as_buffer */
1196 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1197 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001198 islice_doc, /* tp_doc */
1199 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200 0, /* tp_clear */
1201 0, /* tp_richcompare */
1202 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001203 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001204 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001205 0, /* tp_methods */
1206 0, /* tp_members */
1207 0, /* tp_getset */
1208 0, /* tp_base */
1209 0, /* tp_dict */
1210 0, /* tp_descr_get */
1211 0, /* tp_descr_set */
1212 0, /* tp_dictoffset */
1213 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001214 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215 islice_new, /* tp_new */
1216 PyObject_GC_Del, /* tp_free */
1217};
1218
1219
1220/* starmap object ************************************************************/
1221
1222typedef struct {
1223 PyObject_HEAD
1224 PyObject *func;
1225 PyObject *it;
1226} starmapobject;
1227
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001228static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001229
1230static PyObject *
1231starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1232{
1233 PyObject *func, *seq;
1234 PyObject *it;
1235 starmapobject *lz;
1236
1237 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1238 return NULL;
1239
1240 /* Get iterator. */
1241 it = PyObject_GetIter(seq);
1242 if (it == NULL)
1243 return NULL;
1244
1245 /* create starmapobject structure */
1246 lz = (starmapobject *)type->tp_alloc(type, 0);
1247 if (lz == NULL) {
1248 Py_DECREF(it);
1249 return NULL;
1250 }
1251 Py_INCREF(func);
1252 lz->func = func;
1253 lz->it = it;
1254
1255 return (PyObject *)lz;
1256}
1257
1258static void
1259starmap_dealloc(starmapobject *lz)
1260{
1261 PyObject_GC_UnTrack(lz);
1262 Py_XDECREF(lz->func);
1263 Py_XDECREF(lz->it);
1264 lz->ob_type->tp_free(lz);
1265}
1266
1267static int
1268starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1269{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001270 Py_VISIT(lz->it);
1271 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001272 return 0;
1273}
1274
1275static PyObject *
1276starmap_next(starmapobject *lz)
1277{
1278 PyObject *args;
1279 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001280 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001281
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001282 assert(PyIter_Check(it));
1283 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284 if (args == NULL)
1285 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001286 if (!PyTuple_CheckExact(args)) {
1287 Py_DECREF(args);
1288 PyErr_SetString(PyExc_TypeError,
1289 "iterator must return a tuple");
1290 return NULL;
1291 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292 result = PyObject_Call(lz->func, args, NULL);
1293 Py_DECREF(args);
1294 return result;
1295}
1296
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001297PyDoc_STRVAR(starmap_doc,
1298"starmap(function, sequence) --> starmap object\n\
1299\n\
1300Return an iterator whose values are returned from the function evaluated\n\
1301with a argument tuple taken from the given sequence.");
1302
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001303static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304 PyObject_HEAD_INIT(NULL)
1305 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001306 "itertools.starmap", /* tp_name */
1307 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001308 0, /* tp_itemsize */
1309 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001310 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001311 0, /* tp_print */
1312 0, /* tp_getattr */
1313 0, /* tp_setattr */
1314 0, /* tp_compare */
1315 0, /* tp_repr */
1316 0, /* tp_as_number */
1317 0, /* tp_as_sequence */
1318 0, /* tp_as_mapping */
1319 0, /* tp_hash */
1320 0, /* tp_call */
1321 0, /* tp_str */
1322 PyObject_GenericGetAttr, /* tp_getattro */
1323 0, /* tp_setattro */
1324 0, /* tp_as_buffer */
1325 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1326 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001327 starmap_doc, /* tp_doc */
1328 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001329 0, /* tp_clear */
1330 0, /* tp_richcompare */
1331 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001332 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001333 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334 0, /* tp_methods */
1335 0, /* tp_members */
1336 0, /* tp_getset */
1337 0, /* tp_base */
1338 0, /* tp_dict */
1339 0, /* tp_descr_get */
1340 0, /* tp_descr_set */
1341 0, /* tp_dictoffset */
1342 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001343 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001344 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001345 PyObject_GC_Del, /* tp_free */
1346};
1347
1348
1349/* imap object ************************************************************/
1350
1351typedef struct {
1352 PyObject_HEAD
1353 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354 PyObject *func;
1355} imapobject;
1356
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001357static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001358
1359static PyObject *
1360imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1361{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001362 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001363 imapobject *lz;
1364 int numargs, i;
1365
1366 numargs = PyTuple_Size(args);
1367 if (numargs < 2) {
1368 PyErr_SetString(PyExc_TypeError,
1369 "imap() must have at least two arguments.");
1370 return NULL;
1371 }
1372
1373 iters = PyTuple_New(numargs-1);
1374 if (iters == NULL)
1375 return NULL;
1376
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377 for (i=1 ; i<numargs ; i++) {
1378 /* Get iterator. */
1379 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1380 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381 Py_DECREF(iters);
1382 return NULL;
1383 }
1384 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001385 }
1386
1387 /* create imapobject structure */
1388 lz = (imapobject *)type->tp_alloc(type, 0);
1389 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001390 Py_DECREF(iters);
1391 return NULL;
1392 }
1393 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001394 func = PyTuple_GET_ITEM(args, 0);
1395 Py_INCREF(func);
1396 lz->func = func;
1397
1398 return (PyObject *)lz;
1399}
1400
1401static void
1402imap_dealloc(imapobject *lz)
1403{
1404 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001405 Py_XDECREF(lz->iters);
1406 Py_XDECREF(lz->func);
1407 lz->ob_type->tp_free(lz);
1408}
1409
1410static int
1411imap_traverse(imapobject *lz, visitproc visit, void *arg)
1412{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001413 Py_VISIT(lz->iters);
1414 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001415 return 0;
1416}
1417
Raymond Hettinger2012f172003-02-07 05:32:58 +00001418/*
1419imap() is an iterator version of __builtins__.map() except that it does
1420not have the None fill-in feature. That was intentionally left out for
1421the following reasons:
1422
1423 1) Itertools are designed to be easily combined and chained together.
1424 Having all tools stop with the shortest input is a unifying principle
1425 that makes it easier to combine finite iterators (supplying data) with
1426 infinite iterators like count() and repeat() (for supplying sequential
1427 or constant arguments to a function).
1428
1429 2) In typical use cases for combining itertools, having one finite data
1430 supplier run out before another is likely to be an error condition which
1431 should not pass silently by automatically supplying None.
1432
1433 3) The use cases for automatic None fill-in are rare -- not many functions
1434 do something useful when a parameter suddenly switches type and becomes
1435 None.
1436
1437 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001438 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001439
1440 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1441*/
1442
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001443static PyObject *
1444imap_next(imapobject *lz)
1445{
1446 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001447 PyObject *argtuple;
1448 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001449 int numargs, i;
1450
1451 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001452 argtuple = PyTuple_New(numargs);
1453 if (argtuple == NULL)
1454 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001456 for (i=0 ; i<numargs ; i++) {
1457 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1458 if (val == NULL) {
1459 Py_DECREF(argtuple);
1460 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001462 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001464 if (lz->func == Py_None)
1465 return argtuple;
1466 result = PyObject_Call(lz->func, argtuple, NULL);
1467 Py_DECREF(argtuple);
1468 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469}
1470
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471PyDoc_STRVAR(imap_doc,
1472"imap(func, *iterables) --> imap object\n\
1473\n\
1474Make an iterator that computes the function using arguments from\n\
1475each of the iterables. Like map() except that it returns\n\
1476an iterator instead of a list and that it stops when the shortest\n\
1477iterable is exhausted instead of filling in None for shorter\n\
1478iterables.");
1479
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001480static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481 PyObject_HEAD_INIT(NULL)
1482 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001483 "itertools.imap", /* tp_name */
1484 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001485 0, /* tp_itemsize */
1486 /* methods */
1487 (destructor)imap_dealloc, /* tp_dealloc */
1488 0, /* tp_print */
1489 0, /* tp_getattr */
1490 0, /* tp_setattr */
1491 0, /* tp_compare */
1492 0, /* tp_repr */
1493 0, /* tp_as_number */
1494 0, /* tp_as_sequence */
1495 0, /* tp_as_mapping */
1496 0, /* tp_hash */
1497 0, /* tp_call */
1498 0, /* tp_str */
1499 PyObject_GenericGetAttr, /* tp_getattro */
1500 0, /* tp_setattro */
1501 0, /* tp_as_buffer */
1502 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1503 Py_TPFLAGS_BASETYPE, /* tp_flags */
1504 imap_doc, /* tp_doc */
1505 (traverseproc)imap_traverse, /* tp_traverse */
1506 0, /* tp_clear */
1507 0, /* tp_richcompare */
1508 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001509 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001510 (iternextfunc)imap_next, /* tp_iternext */
1511 0, /* tp_methods */
1512 0, /* tp_members */
1513 0, /* tp_getset */
1514 0, /* tp_base */
1515 0, /* tp_dict */
1516 0, /* tp_descr_get */
1517 0, /* tp_descr_set */
1518 0, /* tp_dictoffset */
1519 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001520 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001521 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001522 PyObject_GC_Del, /* tp_free */
1523};
1524
1525
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001526/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001527
1528typedef struct {
1529 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001530 long tuplesize;
1531 long iternum; /* which iterator is active */
1532 PyObject *ittuple; /* tuple of iterators */
1533} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001534
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001535static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536
1537static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001538chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001539{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001540 chainobject *lz;
1541 int tuplesize = PySequence_Length(args);
1542 int i;
1543 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001544
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001545 /* obtain iterators */
1546 assert(PyTuple_Check(args));
1547 ittuple = PyTuple_New(tuplesize);
1548 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001549 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001550 for (i=0; i < tuplesize; ++i) {
1551 PyObject *item = PyTuple_GET_ITEM(args, i);
1552 PyObject *it = PyObject_GetIter(item);
1553 if (it == NULL) {
1554 if (PyErr_ExceptionMatches(PyExc_TypeError))
1555 PyErr_Format(PyExc_TypeError,
1556 "chain argument #%d must support iteration",
1557 i+1);
1558 Py_DECREF(ittuple);
1559 return NULL;
1560 }
1561 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562 }
1563
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001564 /* create chainobject structure */
1565 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001566 if (lz == NULL) {
1567 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001568 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001569 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001570
1571 lz->ittuple = ittuple;
1572 lz->iternum = 0;
1573 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574
1575 return (PyObject *)lz;
1576}
1577
1578static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001579chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580{
1581 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001582 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583 lz->ob_type->tp_free(lz);
1584}
1585
Raymond Hettinger2012f172003-02-07 05:32:58 +00001586static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001587chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001588{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001589 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001590 return 0;
1591}
1592
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001593static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001594chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001596 PyObject *it;
1597 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001598
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001599 while (lz->iternum < lz->tuplesize) {
1600 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1601 item = PyIter_Next(it);
1602 if (item != NULL)
1603 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001604 if (PyErr_Occurred()) {
1605 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1606 PyErr_Clear();
1607 else
1608 return NULL;
1609 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001610 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611 }
1612 return NULL;
1613}
1614
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001615PyDoc_STRVAR(chain_doc,
1616"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001618Return a chain object whose .next() method returns elements from the\n\
1619first iterable until it is exhausted, then elements from the next\n\
1620iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001621
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001622static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001623 PyObject_HEAD_INIT(NULL)
1624 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001625 "itertools.chain", /* tp_name */
1626 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001627 0, /* tp_itemsize */
1628 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001629 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001630 0, /* tp_print */
1631 0, /* tp_getattr */
1632 0, /* tp_setattr */
1633 0, /* tp_compare */
1634 0, /* tp_repr */
1635 0, /* tp_as_number */
1636 0, /* tp_as_sequence */
1637 0, /* tp_as_mapping */
1638 0, /* tp_hash */
1639 0, /* tp_call */
1640 0, /* tp_str */
1641 PyObject_GenericGetAttr, /* tp_getattro */
1642 0, /* tp_setattro */
1643 0, /* tp_as_buffer */
1644 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1645 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001646 chain_doc, /* tp_doc */
1647 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648 0, /* tp_clear */
1649 0, /* tp_richcompare */
1650 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001651 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001652 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001653 0, /* tp_methods */
1654 0, /* tp_members */
1655 0, /* tp_getset */
1656 0, /* tp_base */
1657 0, /* tp_dict */
1658 0, /* tp_descr_get */
1659 0, /* tp_descr_set */
1660 0, /* tp_dictoffset */
1661 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001662 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001663 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001664 PyObject_GC_Del, /* tp_free */
1665};
1666
1667
1668/* ifilter object ************************************************************/
1669
1670typedef struct {
1671 PyObject_HEAD
1672 PyObject *func;
1673 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674} ifilterobject;
1675
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001676static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677
1678static PyObject *
1679ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1680{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001681 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001682 PyObject *it;
1683 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001684
Raymond Hettinger60eca932003-02-09 06:40:58 +00001685 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001686 return NULL;
1687
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001688 /* Get iterator. */
1689 it = PyObject_GetIter(seq);
1690 if (it == NULL)
1691 return NULL;
1692
1693 /* create ifilterobject structure */
1694 lz = (ifilterobject *)type->tp_alloc(type, 0);
1695 if (lz == NULL) {
1696 Py_DECREF(it);
1697 return NULL;
1698 }
1699 Py_INCREF(func);
1700 lz->func = func;
1701 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001702
1703 return (PyObject *)lz;
1704}
1705
1706static void
1707ifilter_dealloc(ifilterobject *lz)
1708{
1709 PyObject_GC_UnTrack(lz);
1710 Py_XDECREF(lz->func);
1711 Py_XDECREF(lz->it);
1712 lz->ob_type->tp_free(lz);
1713}
1714
1715static int
1716ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1717{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001718 Py_VISIT(lz->it);
1719 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001720 return 0;
1721}
1722
1723static PyObject *
1724ifilter_next(ifilterobject *lz)
1725{
1726 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001727 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001728 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001729 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001730
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001731 assert(PyIter_Check(it));
1732 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001733 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001734 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735 if (item == NULL)
1736 return NULL;
1737
1738 if (lz->func == Py_None) {
1739 ok = PyObject_IsTrue(item);
1740 } else {
1741 PyObject *good;
1742 good = PyObject_CallFunctionObjArgs(lz->func,
1743 item, NULL);
1744 if (good == NULL) {
1745 Py_DECREF(item);
1746 return NULL;
1747 }
1748 ok = PyObject_IsTrue(good);
1749 Py_DECREF(good);
1750 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001751 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001752 return item;
1753 Py_DECREF(item);
1754 }
1755}
1756
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001758"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001760Return those items of sequence for which function(item) is true.\n\
1761If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001762
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001763static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764 PyObject_HEAD_INIT(NULL)
1765 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001766 "itertools.ifilter", /* tp_name */
1767 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001768 0, /* tp_itemsize */
1769 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001770 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001771 0, /* tp_print */
1772 0, /* tp_getattr */
1773 0, /* tp_setattr */
1774 0, /* tp_compare */
1775 0, /* tp_repr */
1776 0, /* tp_as_number */
1777 0, /* tp_as_sequence */
1778 0, /* tp_as_mapping */
1779 0, /* tp_hash */
1780 0, /* tp_call */
1781 0, /* tp_str */
1782 PyObject_GenericGetAttr, /* tp_getattro */
1783 0, /* tp_setattro */
1784 0, /* tp_as_buffer */
1785 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1786 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001787 ifilter_doc, /* tp_doc */
1788 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001789 0, /* tp_clear */
1790 0, /* tp_richcompare */
1791 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001792 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001793 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001794 0, /* tp_methods */
1795 0, /* tp_members */
1796 0, /* tp_getset */
1797 0, /* tp_base */
1798 0, /* tp_dict */
1799 0, /* tp_descr_get */
1800 0, /* tp_descr_set */
1801 0, /* tp_dictoffset */
1802 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001803 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001804 ifilter_new, /* tp_new */
1805 PyObject_GC_Del, /* tp_free */
1806};
1807
1808
1809/* ifilterfalse object ************************************************************/
1810
1811typedef struct {
1812 PyObject_HEAD
1813 PyObject *func;
1814 PyObject *it;
1815} ifilterfalseobject;
1816
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001817static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001818
1819static PyObject *
1820ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1821{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001822 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001823 PyObject *it;
1824 ifilterfalseobject *lz;
1825
1826 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1827 return NULL;
1828
1829 /* Get iterator. */
1830 it = PyObject_GetIter(seq);
1831 if (it == NULL)
1832 return NULL;
1833
1834 /* create ifilterfalseobject structure */
1835 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1836 if (lz == NULL) {
1837 Py_DECREF(it);
1838 return NULL;
1839 }
1840 Py_INCREF(func);
1841 lz->func = func;
1842 lz->it = it;
1843
1844 return (PyObject *)lz;
1845}
1846
1847static void
1848ifilterfalse_dealloc(ifilterfalseobject *lz)
1849{
1850 PyObject_GC_UnTrack(lz);
1851 Py_XDECREF(lz->func);
1852 Py_XDECREF(lz->it);
1853 lz->ob_type->tp_free(lz);
1854}
1855
1856static int
1857ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1858{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001859 Py_VISIT(lz->it);
1860 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001861 return 0;
1862}
1863
1864static PyObject *
1865ifilterfalse_next(ifilterfalseobject *lz)
1866{
1867 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001868 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001869 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001870 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001871
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001872 assert(PyIter_Check(it));
1873 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001874 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001875 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001876 if (item == NULL)
1877 return NULL;
1878
1879 if (lz->func == Py_None) {
1880 ok = PyObject_IsTrue(item);
1881 } else {
1882 PyObject *good;
1883 good = PyObject_CallFunctionObjArgs(lz->func,
1884 item, NULL);
1885 if (good == NULL) {
1886 Py_DECREF(item);
1887 return NULL;
1888 }
1889 ok = PyObject_IsTrue(good);
1890 Py_DECREF(good);
1891 }
1892 if (!ok)
1893 return item;
1894 Py_DECREF(item);
1895 }
1896}
1897
Raymond Hettinger60eca932003-02-09 06:40:58 +00001898PyDoc_STRVAR(ifilterfalse_doc,
1899"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1900\n\
1901Return those items of sequence for which function(item) is false.\n\
1902If function is None, return the items that are false.");
1903
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001904static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001905 PyObject_HEAD_INIT(NULL)
1906 0, /* ob_size */
1907 "itertools.ifilterfalse", /* tp_name */
1908 sizeof(ifilterfalseobject), /* tp_basicsize */
1909 0, /* tp_itemsize */
1910 /* methods */
1911 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1912 0, /* tp_print */
1913 0, /* tp_getattr */
1914 0, /* tp_setattr */
1915 0, /* tp_compare */
1916 0, /* tp_repr */
1917 0, /* tp_as_number */
1918 0, /* tp_as_sequence */
1919 0, /* tp_as_mapping */
1920 0, /* tp_hash */
1921 0, /* tp_call */
1922 0, /* tp_str */
1923 PyObject_GenericGetAttr, /* tp_getattro */
1924 0, /* tp_setattro */
1925 0, /* tp_as_buffer */
1926 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1927 Py_TPFLAGS_BASETYPE, /* tp_flags */
1928 ifilterfalse_doc, /* tp_doc */
1929 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1930 0, /* tp_clear */
1931 0, /* tp_richcompare */
1932 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001933 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001934 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1935 0, /* tp_methods */
1936 0, /* tp_members */
1937 0, /* tp_getset */
1938 0, /* tp_base */
1939 0, /* tp_dict */
1940 0, /* tp_descr_get */
1941 0, /* tp_descr_set */
1942 0, /* tp_dictoffset */
1943 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001944 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001945 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001946 PyObject_GC_Del, /* tp_free */
1947};
1948
1949
1950/* count object ************************************************************/
1951
1952typedef struct {
1953 PyObject_HEAD
1954 long cnt;
1955} countobject;
1956
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001957static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001958
1959static PyObject *
1960count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1961{
1962 countobject *lz;
1963 long cnt = 0;
1964
1965 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1966 return NULL;
1967
1968 /* create countobject structure */
1969 lz = (countobject *)PyObject_New(countobject, &count_type);
1970 if (lz == NULL)
1971 return NULL;
1972 lz->cnt = cnt;
1973
1974 return (PyObject *)lz;
1975}
1976
1977static PyObject *
1978count_next(countobject *lz)
1979{
1980 return PyInt_FromLong(lz->cnt++);
1981}
1982
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001983static PyObject *
1984count_repr(countobject *lz)
1985{
Brett Cannon00468392004-04-13 02:43:53 +00001986 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001987}
1988
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001989PyDoc_STRVAR(count_doc,
1990"count([firstval]) --> count object\n\
1991\n\
1992Return a count object whose .next() method returns consecutive\n\
1993integers starting from zero or, if specified, from firstval.");
1994
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001995static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001996 PyObject_HEAD_INIT(NULL)
1997 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001998 "itertools.count", /* tp_name */
1999 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002000 0, /* tp_itemsize */
2001 /* methods */
2002 (destructor)PyObject_Del, /* tp_dealloc */
2003 0, /* tp_print */
2004 0, /* tp_getattr */
2005 0, /* tp_setattr */
2006 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002007 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002008 0, /* tp_as_number */
2009 0, /* tp_as_sequence */
2010 0, /* tp_as_mapping */
2011 0, /* tp_hash */
2012 0, /* tp_call */
2013 0, /* tp_str */
2014 PyObject_GenericGetAttr, /* tp_getattro */
2015 0, /* tp_setattro */
2016 0, /* tp_as_buffer */
2017 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002018 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002019 0, /* tp_traverse */
2020 0, /* tp_clear */
2021 0, /* tp_richcompare */
2022 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002023 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002024 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002025 0, /* tp_methods */
2026 0, /* tp_members */
2027 0, /* tp_getset */
2028 0, /* tp_base */
2029 0, /* tp_dict */
2030 0, /* tp_descr_get */
2031 0, /* tp_descr_set */
2032 0, /* tp_dictoffset */
2033 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002034 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002035 count_new, /* tp_new */
2036};
2037
2038
2039/* izip object ************************************************************/
2040
2041#include "Python.h"
2042
2043typedef struct {
2044 PyObject_HEAD
2045 long tuplesize;
2046 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002047 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002048} izipobject;
2049
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002050static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002051
2052static PyObject *
2053izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2054{
2055 izipobject *lz;
2056 int i;
2057 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002058 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002059 int tuplesize = PySequence_Length(args);
2060
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002061 /* args must be a tuple */
2062 assert(PyTuple_Check(args));
2063
2064 /* obtain iterators */
2065 ittuple = PyTuple_New(tuplesize);
2066 if(ittuple == NULL)
2067 return NULL;
2068 for (i=0; i < tuplesize; ++i) {
2069 PyObject *item = PyTuple_GET_ITEM(args, i);
2070 PyObject *it = PyObject_GetIter(item);
2071 if (it == NULL) {
2072 if (PyErr_ExceptionMatches(PyExc_TypeError))
2073 PyErr_Format(PyExc_TypeError,
2074 "izip argument #%d must support iteration",
2075 i+1);
2076 Py_DECREF(ittuple);
2077 return NULL;
2078 }
2079 PyTuple_SET_ITEM(ittuple, i, it);
2080 }
2081
Raymond Hettinger2012f172003-02-07 05:32:58 +00002082 /* create a result holder */
2083 result = PyTuple_New(tuplesize);
2084 if (result == NULL) {
2085 Py_DECREF(ittuple);
2086 return NULL;
2087 }
2088 for (i=0 ; i < tuplesize ; i++) {
2089 Py_INCREF(Py_None);
2090 PyTuple_SET_ITEM(result, i, Py_None);
2091 }
2092
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002093 /* create izipobject structure */
2094 lz = (izipobject *)type->tp_alloc(type, 0);
2095 if (lz == NULL) {
2096 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002097 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002098 return NULL;
2099 }
2100 lz->ittuple = ittuple;
2101 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002102 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002103
2104 return (PyObject *)lz;
2105}
2106
2107static void
2108izip_dealloc(izipobject *lz)
2109{
2110 PyObject_GC_UnTrack(lz);
2111 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002112 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002113 lz->ob_type->tp_free(lz);
2114}
2115
2116static int
2117izip_traverse(izipobject *lz, visitproc visit, void *arg)
2118{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002119 Py_VISIT(lz->ittuple);
2120 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002121 return 0;
2122}
2123
2124static PyObject *
2125izip_next(izipobject *lz)
2126{
2127 int i;
2128 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002129 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002130 PyObject *it;
2131 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002132 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002133
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002134 if (tuplesize == 0)
2135 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002136 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002137 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002138 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002139 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002140 assert(PyIter_Check(it));
2141 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002142 if (item == NULL) {
2143 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002144 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002145 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002146 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002147 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002148 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002149 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002150 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002151 result = PyTuple_New(tuplesize);
2152 if (result == NULL)
2153 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002154 for (i=0 ; i < tuplesize ; i++) {
2155 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002156 assert(PyIter_Check(it));
2157 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002158 if (item == NULL) {
2159 Py_DECREF(result);
2160 return NULL;
2161 }
2162 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002163 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002164 }
2165 return result;
2166}
2167
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002168PyDoc_STRVAR(izip_doc,
2169"izip(iter1 [,iter2 [...]]) --> izip object\n\
2170\n\
2171Return a izip object whose .next() method returns a tuple where\n\
2172the i-th element comes from the i-th iterable argument. The .next()\n\
2173method continues until the shortest iterable in the argument sequence\n\
2174is exhausted and then it raises StopIteration. Works like the zip()\n\
2175function but consumes less memory by returning an iterator instead of\n\
2176a list.");
2177
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002178static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002179 PyObject_HEAD_INIT(NULL)
2180 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002181 "itertools.izip", /* tp_name */
2182 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002183 0, /* tp_itemsize */
2184 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002185 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002186 0, /* tp_print */
2187 0, /* tp_getattr */
2188 0, /* tp_setattr */
2189 0, /* tp_compare */
2190 0, /* tp_repr */
2191 0, /* tp_as_number */
2192 0, /* tp_as_sequence */
2193 0, /* tp_as_mapping */
2194 0, /* tp_hash */
2195 0, /* tp_call */
2196 0, /* tp_str */
2197 PyObject_GenericGetAttr, /* tp_getattro */
2198 0, /* tp_setattro */
2199 0, /* tp_as_buffer */
2200 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2201 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002202 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002203 (traverseproc)izip_traverse, /* tp_traverse */
2204 0, /* tp_clear */
2205 0, /* tp_richcompare */
2206 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002207 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002208 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002209 0, /* tp_methods */
2210 0, /* tp_members */
2211 0, /* tp_getset */
2212 0, /* tp_base */
2213 0, /* tp_dict */
2214 0, /* tp_descr_get */
2215 0, /* tp_descr_set */
2216 0, /* tp_dictoffset */
2217 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002218 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002219 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002220 PyObject_GC_Del, /* tp_free */
2221};
2222
2223
2224/* repeat object ************************************************************/
2225
2226typedef struct {
2227 PyObject_HEAD
2228 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002229 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002230} repeatobject;
2231
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002232static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002233
2234static PyObject *
2235repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2236{
2237 repeatobject *ro;
2238 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002239 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002240
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002241 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002242 return NULL;
2243
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002244 if (PyTuple_Size(args) == 2 && cnt < 0)
2245 cnt = 0;
2246
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002247 ro = (repeatobject *)type->tp_alloc(type, 0);
2248 if (ro == NULL)
2249 return NULL;
2250 Py_INCREF(element);
2251 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002252 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002253 return (PyObject *)ro;
2254}
2255
2256static void
2257repeat_dealloc(repeatobject *ro)
2258{
2259 PyObject_GC_UnTrack(ro);
2260 Py_XDECREF(ro->element);
2261 ro->ob_type->tp_free(ro);
2262}
2263
2264static int
2265repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2266{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002267 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002268 return 0;
2269}
2270
2271static PyObject *
2272repeat_next(repeatobject *ro)
2273{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002274 if (ro->cnt == 0)
2275 return NULL;
2276 if (ro->cnt > 0)
2277 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002278 Py_INCREF(ro->element);
2279 return ro->element;
2280}
2281
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002282static PyObject *
2283repeat_repr(repeatobject *ro)
2284{
2285 PyObject *result, *objrepr;
2286
2287 objrepr = PyObject_Repr(ro->element);
2288 if (objrepr == NULL)
2289 return NULL;
2290
2291 if (ro->cnt == -1)
2292 result = PyString_FromFormat("repeat(%s)",
2293 PyString_AS_STRING(objrepr));
2294 else
Brett Cannon00468392004-04-13 02:43:53 +00002295 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002296 PyString_AS_STRING(objrepr), ro->cnt);
2297 Py_DECREF(objrepr);
2298 return result;
2299}
2300
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002301static int
2302repeat_len(repeatobject *ro)
2303{
2304 if (ro->cnt == -1)
2305 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2306 return (int)(ro->cnt);
2307}
2308
2309static PySequenceMethods repeat_as_sequence = {
2310 (inquiry)repeat_len, /* sq_length */
2311 0, /* sq_concat */
2312};
2313
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002314PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002315"repeat(element [,times]) -> create an iterator which returns the element\n\
2316for the specified number of times. If not specified, returns the element\n\
2317endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002318
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002319static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002320 PyObject_HEAD_INIT(NULL)
2321 0, /* ob_size */
2322 "itertools.repeat", /* tp_name */
2323 sizeof(repeatobject), /* tp_basicsize */
2324 0, /* tp_itemsize */
2325 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002326 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002327 0, /* tp_print */
2328 0, /* tp_getattr */
2329 0, /* tp_setattr */
2330 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002331 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002332 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002333 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002334 0, /* tp_as_mapping */
2335 0, /* tp_hash */
2336 0, /* tp_call */
2337 0, /* tp_str */
2338 PyObject_GenericGetAttr, /* tp_getattro */
2339 0, /* tp_setattro */
2340 0, /* tp_as_buffer */
2341 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2342 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002343 repeat_doc, /* tp_doc */
2344 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002345 0, /* tp_clear */
2346 0, /* tp_richcompare */
2347 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002348 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002349 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002350 0, /* tp_methods */
2351 0, /* tp_members */
2352 0, /* tp_getset */
2353 0, /* tp_base */
2354 0, /* tp_dict */
2355 0, /* tp_descr_get */
2356 0, /* tp_descr_set */
2357 0, /* tp_dictoffset */
2358 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002359 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002360 repeat_new, /* tp_new */
2361 PyObject_GC_Del, /* tp_free */
2362};
2363
2364
2365/* module level code ********************************************************/
2366
2367PyDoc_STRVAR(module_doc,
2368"Functional tools for creating and using iterators.\n\
2369\n\
2370Infinite iterators:\n\
2371count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002372cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002373repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002374\n\
2375Iterators terminating on the shortest input sequence:\n\
2376izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002377ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2378ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002379islice(seq, [start,] stop [, step]) --> elements from\n\
2380 seq[start:stop:step]\n\
2381imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2382starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002383tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002384chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002385takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2386dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002387groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002388");
2389
2390
Raymond Hettingerad983e72003-11-12 14:32:26 +00002391static PyMethodDef module_methods[] = {
2392 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2393 {NULL, NULL} /* sentinel */
2394};
2395
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002396PyMODINIT_FUNC
2397inititertools(void)
2398{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002399 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002400 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002401 char *name;
2402 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002403 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002404 &dropwhile_type,
2405 &takewhile_type,
2406 &islice_type,
2407 &starmap_type,
2408 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002409 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002410 &ifilter_type,
2411 &ifilterfalse_type,
2412 &count_type,
2413 &izip_type,
2414 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002415 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002416 NULL
2417 };
2418
Skip Montanarof3938fd2004-02-10 20:27:40 +00002419 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002420 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002421
Raymond Hettinger60eca932003-02-09 06:40:58 +00002422 for (i=0 ; typelist[i] != NULL ; i++) {
2423 if (PyType_Ready(typelist[i]) < 0)
2424 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002425 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002426 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002427 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002428 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002429 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002430
2431 if (PyType_Ready(&teedataobject_type) < 0)
2432 return;
2433 if (PyType_Ready(&tee_type) < 0)
2434 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002435 if (PyType_Ready(&_grouper_type) < 0)
2436 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002437}