blob: 71081fbe43b6f13b84265bbe75d4f0753efd4a03 [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{
Martin v. Löwis02cbf4a2006-02-27 17:20:04 +000029 static char *kwargs[] = {"iterable", "key", NULL};
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030 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
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000343 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344 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;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000351 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000352 return (PyObject *)tdo;
353}
354
355static PyObject *
356teedataobject_jumplink(teedataobject *tdo)
357{
358 if (tdo->nextlink == NULL)
359 tdo->nextlink = teedataobject_new(tdo->it);
360 Py_INCREF(tdo->nextlink);
361 return tdo->nextlink;
362}
363
364static PyObject *
365teedataobject_getitem(teedataobject *tdo, int i)
366{
367 PyObject *value;
368
369 assert(i < LINKCELLS);
370 if (i < tdo->numread)
371 value = tdo->values[i];
372 else {
373 /* this is the lead iterator, so fetch more data */
374 assert(i == tdo->numread);
375 value = PyIter_Next(tdo->it);
376 if (value == NULL)
377 return NULL;
378 tdo->numread++;
379 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000380 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000381 Py_INCREF(value);
382 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000383}
384
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000385static int
386teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
387{
388 int i;
389 Py_VISIT(tdo->it);
390 for (i = 0; i < tdo->numread; i++)
391 Py_VISIT(tdo->values[i]);
392 Py_VISIT(tdo->nextlink);
393 return 0;
394}
395
396static int
397teedataobject_clear(teedataobject *tdo)
398{
399 int i;
400 Py_CLEAR(tdo->it);
401 for (i=0 ; i<tdo->numread ; i++)
402 Py_CLEAR(tdo->values[i]);
403 Py_CLEAR(tdo->nextlink);
404 return 0;
405}
406
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000407static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000408teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000409{
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000410 PyObject_GC_UnTrack(tdo);
411 teedataobject_clear(tdo);
412 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000413}
414
Raymond Hettingerad983e72003-11-12 14:32:26 +0000415PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000416
Raymond Hettingerad983e72003-11-12 14:32:26 +0000417static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000418 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000419 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000420 "itertools.tee_dataobject", /* tp_name */
421 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422 0, /* tp_itemsize */
423 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000424 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000425 0, /* tp_print */
426 0, /* tp_getattr */
427 0, /* tp_setattr */
428 0, /* tp_compare */
429 0, /* tp_repr */
430 0, /* tp_as_number */
431 0, /* tp_as_sequence */
432 0, /* tp_as_mapping */
433 0, /* tp_hash */
434 0, /* tp_call */
435 0, /* tp_str */
436 PyObject_GenericGetAttr, /* tp_getattro */
437 0, /* tp_setattro */
438 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000439 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000440 teedataobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000441 (traverseproc)teedataobject_traverse, /* tp_traverse */
442 (inquiry)teedataobject_clear, /* tp_clear */
443 0, /* tp_richcompare */
444 0, /* tp_weaklistoffset */
445 0, /* tp_iter */
446 0, /* tp_iternext */
447 0, /* tp_methods */
448 0, /* tp_members */
449 0, /* tp_getset */
450 0, /* tp_base */
451 0, /* tp_dict */
452 0, /* tp_descr_get */
453 0, /* tp_descr_set */
454 0, /* tp_dictoffset */
455 0, /* tp_init */
456 0, /* tp_alloc */
457 0, /* tp_new */
458 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000459};
460
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000461
462static PyTypeObject tee_type;
463
464static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000465tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000466{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000467 PyObject *value, *link;
468
469 if (to->index >= LINKCELLS) {
470 link = teedataobject_jumplink(to->dataobj);
471 Py_XDECREF(to->dataobj);
472 to->dataobj = (teedataobject *)link;
473 to->index = 0;
474 }
475 value = teedataobject_getitem(to->dataobj, to->index);
476 if (value == NULL)
477 return NULL;
478 to->index++;
479 return value;
480}
481
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000482static int
483tee_traverse(teeobject *to, visitproc visit, void *arg)
484{
485 Py_VISIT((PyObject *)to->dataobj);
486 return 0;
487}
488
Raymond Hettingerad983e72003-11-12 14:32:26 +0000489static PyObject *
490tee_copy(teeobject *to)
491{
492 teeobject *newto;
493
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000494 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000495 if (newto == NULL)
496 return NULL;
497 Py_INCREF(to->dataobj);
498 newto->dataobj = to->dataobj;
499 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000500 newto->weakreflist = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000501 return (PyObject *)newto;
502}
503
504PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
505
506static PyObject *
507tee_fromiterable(PyObject *iterable)
508{
509 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000510 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000511
512 it = PyObject_GetIter(iterable);
513 if (it == NULL)
514 return NULL;
515 if (PyObject_TypeCheck(it, &tee_type)) {
516 to = (teeobject *)tee_copy((teeobject *)it);
517 goto done;
518 }
519
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000520 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000521 if (to == NULL)
522 goto done;
523 to->dataobj = (teedataobject *)teedataobject_new(it);
524 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000525 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000526 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000527done:
528 Py_XDECREF(it);
529 return (PyObject *)to;
530}
531
532static PyObject *
533tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
534{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000535 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000536
537 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
538 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000539 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000540}
541
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000542static int
543tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000544{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000545 if (to->weakreflist != NULL)
546 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000547 Py_CLEAR(to->dataobj);
548 return 0;
549}
550
551static void
552tee_dealloc(teeobject *to)
553{
554 PyObject_GC_UnTrack(to);
555 tee_clear(to);
556 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000557}
558
Raymond Hettingerad983e72003-11-12 14:32:26 +0000559PyDoc_STRVAR(teeobject_doc,
560"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000561
Raymond Hettingerad983e72003-11-12 14:32:26 +0000562static PyMethodDef tee_methods[] = {
563 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
564 {NULL, NULL} /* sentinel */
565};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000566
567static PyTypeObject tee_type = {
568 PyObject_HEAD_INIT(NULL)
569 0, /* ob_size */
570 "itertools.tee", /* tp_name */
571 sizeof(teeobject), /* tp_basicsize */
572 0, /* tp_itemsize */
573 /* methods */
574 (destructor)tee_dealloc, /* tp_dealloc */
575 0, /* tp_print */
576 0, /* tp_getattr */
577 0, /* tp_setattr */
578 0, /* tp_compare */
579 0, /* tp_repr */
580 0, /* tp_as_number */
581 0, /* tp_as_sequence */
582 0, /* tp_as_mapping */
583 0, /* tp_hash */
584 0, /* tp_call */
585 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000586 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000587 0, /* tp_setattro */
588 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000589 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000590 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000591 (traverseproc)tee_traverse, /* tp_traverse */
592 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000593 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000594 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000595 PyObject_SelfIter, /* tp_iter */
596 (iternextfunc)tee_next, /* tp_iternext */
597 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000598 0, /* tp_members */
599 0, /* tp_getset */
600 0, /* tp_base */
601 0, /* tp_dict */
602 0, /* tp_descr_get */
603 0, /* tp_descr_set */
604 0, /* tp_dictoffset */
605 0, /* tp_init */
606 0, /* tp_alloc */
607 tee_new, /* tp_new */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000608 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000609};
610
Raymond Hettingerad983e72003-11-12 14:32:26 +0000611static PyObject *
612tee(PyObject *self, PyObject *args)
613{
614 int i, n=2;
615 PyObject *it, *iterable, *copyable, *result;
616
617 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
618 return NULL;
619 result = PyTuple_New(n);
620 if (result == NULL)
621 return NULL;
622 if (n == 0)
623 return result;
624 it = PyObject_GetIter(iterable);
625 if (it == NULL) {
626 Py_DECREF(result);
627 return NULL;
628 }
629 if (!PyObject_HasAttrString(it, "__copy__")) {
630 copyable = tee_fromiterable(it);
631 Py_DECREF(it);
632 if (copyable == NULL) {
633 Py_DECREF(result);
634 return NULL;
635 }
636 } else
637 copyable = it;
638 PyTuple_SET_ITEM(result, 0, copyable);
639 for (i=1 ; i<n ; i++) {
640 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
641 if (copyable == NULL) {
642 Py_DECREF(result);
643 return NULL;
644 }
645 PyTuple_SET_ITEM(result, i, copyable);
646 }
647 return result;
648}
649
650PyDoc_STRVAR(tee_doc,
651"tee(iterable, n=2) --> tuple of n independent iterators.");
652
653
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000654/* cycle object **********************************************************/
655
656typedef struct {
657 PyObject_HEAD
658 PyObject *it;
659 PyObject *saved;
660 int firstpass;
661} cycleobject;
662
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000663static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000664
665static PyObject *
666cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
667{
668 PyObject *it;
669 PyObject *iterable;
670 PyObject *saved;
671 cycleobject *lz;
672
Georg Brandl02c42872005-08-26 06:42:30 +0000673 if (!_PyArg_NoKeywords("cycle()", kwds))
674 return NULL;
675
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000676 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
677 return NULL;
678
679 /* Get iterator. */
680 it = PyObject_GetIter(iterable);
681 if (it == NULL)
682 return NULL;
683
684 saved = PyList_New(0);
685 if (saved == NULL) {
686 Py_DECREF(it);
687 return NULL;
688 }
689
690 /* create cycleobject structure */
691 lz = (cycleobject *)type->tp_alloc(type, 0);
692 if (lz == NULL) {
693 Py_DECREF(it);
694 Py_DECREF(saved);
695 return NULL;
696 }
697 lz->it = it;
698 lz->saved = saved;
699 lz->firstpass = 0;
700
701 return (PyObject *)lz;
702}
703
704static void
705cycle_dealloc(cycleobject *lz)
706{
707 PyObject_GC_UnTrack(lz);
708 Py_XDECREF(lz->saved);
709 Py_XDECREF(lz->it);
710 lz->ob_type->tp_free(lz);
711}
712
713static int
714cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
715{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000716 Py_VISIT(lz->it);
717 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000718 return 0;
719}
720
721static PyObject *
722cycle_next(cycleobject *lz)
723{
724 PyObject *item;
725 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000726 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000727
728 while (1) {
729 item = PyIter_Next(lz->it);
730 if (item != NULL) {
731 if (!lz->firstpass)
732 PyList_Append(lz->saved, item);
733 return item;
734 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000735 if (PyErr_Occurred()) {
736 if (PyErr_ExceptionMatches(PyExc_StopIteration))
737 PyErr_Clear();
738 else
739 return NULL;
740 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000741 if (PyList_Size(lz->saved) == 0)
742 return NULL;
743 it = PyObject_GetIter(lz->saved);
744 if (it == NULL)
745 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000746 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000747 lz->it = it;
748 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000749 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000750 }
751}
752
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000753PyDoc_STRVAR(cycle_doc,
754"cycle(iterable) --> cycle object\n\
755\n\
756Return elements from the iterable until it is exhausted.\n\
757Then repeat the sequence indefinitely.");
758
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000759static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760 PyObject_HEAD_INIT(NULL)
761 0, /* ob_size */
762 "itertools.cycle", /* tp_name */
763 sizeof(cycleobject), /* tp_basicsize */
764 0, /* tp_itemsize */
765 /* methods */
766 (destructor)cycle_dealloc, /* tp_dealloc */
767 0, /* tp_print */
768 0, /* tp_getattr */
769 0, /* tp_setattr */
770 0, /* tp_compare */
771 0, /* tp_repr */
772 0, /* tp_as_number */
773 0, /* tp_as_sequence */
774 0, /* tp_as_mapping */
775 0, /* tp_hash */
776 0, /* tp_call */
777 0, /* tp_str */
778 PyObject_GenericGetAttr, /* tp_getattro */
779 0, /* tp_setattro */
780 0, /* tp_as_buffer */
781 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
782 Py_TPFLAGS_BASETYPE, /* tp_flags */
783 cycle_doc, /* tp_doc */
784 (traverseproc)cycle_traverse, /* tp_traverse */
785 0, /* tp_clear */
786 0, /* tp_richcompare */
787 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000788 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000789 (iternextfunc)cycle_next, /* tp_iternext */
790 0, /* tp_methods */
791 0, /* tp_members */
792 0, /* tp_getset */
793 0, /* tp_base */
794 0, /* tp_dict */
795 0, /* tp_descr_get */
796 0, /* tp_descr_set */
797 0, /* tp_dictoffset */
798 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000799 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000800 cycle_new, /* tp_new */
801 PyObject_GC_Del, /* tp_free */
802};
803
804
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000805/* dropwhile object **********************************************************/
806
807typedef struct {
808 PyObject_HEAD
809 PyObject *func;
810 PyObject *it;
811 long start;
812} dropwhileobject;
813
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000814static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000815
816static PyObject *
817dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
818{
819 PyObject *func, *seq;
820 PyObject *it;
821 dropwhileobject *lz;
822
Georg Brandl02c42872005-08-26 06:42:30 +0000823 if (!_PyArg_NoKeywords("dropwhile()", kwds))
824 return NULL;
825
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000826 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
827 return NULL;
828
829 /* Get iterator. */
830 it = PyObject_GetIter(seq);
831 if (it == NULL)
832 return NULL;
833
834 /* create dropwhileobject structure */
835 lz = (dropwhileobject *)type->tp_alloc(type, 0);
836 if (lz == NULL) {
837 Py_DECREF(it);
838 return NULL;
839 }
840 Py_INCREF(func);
841 lz->func = func;
842 lz->it = it;
843 lz->start = 0;
844
845 return (PyObject *)lz;
846}
847
848static void
849dropwhile_dealloc(dropwhileobject *lz)
850{
851 PyObject_GC_UnTrack(lz);
852 Py_XDECREF(lz->func);
853 Py_XDECREF(lz->it);
854 lz->ob_type->tp_free(lz);
855}
856
857static int
858dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
859{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000860 Py_VISIT(lz->it);
861 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862 return 0;
863}
864
865static PyObject *
866dropwhile_next(dropwhileobject *lz)
867{
868 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000869 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000870 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000871 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000872
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000873 assert(PyIter_Check(it));
874 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000875 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000876 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000877 if (item == NULL)
878 return NULL;
879 if (lz->start == 1)
880 return item;
881
882 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
883 if (good == NULL) {
884 Py_DECREF(item);
885 return NULL;
886 }
887 ok = PyObject_IsTrue(good);
888 Py_DECREF(good);
889 if (!ok) {
890 lz->start = 1;
891 return item;
892 }
893 Py_DECREF(item);
894 }
895}
896
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000897PyDoc_STRVAR(dropwhile_doc,
898"dropwhile(predicate, iterable) --> dropwhile object\n\
899\n\
900Drop items from the iterable while predicate(item) is true.\n\
901Afterwards, return every element until the iterable is exhausted.");
902
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000903static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000904 PyObject_HEAD_INIT(NULL)
905 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000906 "itertools.dropwhile", /* tp_name */
907 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000908 0, /* tp_itemsize */
909 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000910 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000911 0, /* tp_print */
912 0, /* tp_getattr */
913 0, /* tp_setattr */
914 0, /* tp_compare */
915 0, /* tp_repr */
916 0, /* tp_as_number */
917 0, /* tp_as_sequence */
918 0, /* tp_as_mapping */
919 0, /* tp_hash */
920 0, /* tp_call */
921 0, /* tp_str */
922 PyObject_GenericGetAttr, /* tp_getattro */
923 0, /* tp_setattro */
924 0, /* tp_as_buffer */
925 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
926 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000927 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000928 (traverseproc)dropwhile_traverse, /* tp_traverse */
929 0, /* tp_clear */
930 0, /* tp_richcompare */
931 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000932 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000933 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000934 0, /* tp_methods */
935 0, /* tp_members */
936 0, /* tp_getset */
937 0, /* tp_base */
938 0, /* tp_dict */
939 0, /* tp_descr_get */
940 0, /* tp_descr_set */
941 0, /* tp_dictoffset */
942 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000943 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000944 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000945 PyObject_GC_Del, /* tp_free */
946};
947
948
949/* takewhile object **********************************************************/
950
951typedef struct {
952 PyObject_HEAD
953 PyObject *func;
954 PyObject *it;
955 long stop;
956} takewhileobject;
957
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000958static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000959
960static PyObject *
961takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
962{
963 PyObject *func, *seq;
964 PyObject *it;
965 takewhileobject *lz;
966
Georg Brandl02c42872005-08-26 06:42:30 +0000967 if (!_PyArg_NoKeywords("takewhile()", kwds))
968 return NULL;
969
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000970 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
971 return NULL;
972
973 /* Get iterator. */
974 it = PyObject_GetIter(seq);
975 if (it == NULL)
976 return NULL;
977
978 /* create takewhileobject structure */
979 lz = (takewhileobject *)type->tp_alloc(type, 0);
980 if (lz == NULL) {
981 Py_DECREF(it);
982 return NULL;
983 }
984 Py_INCREF(func);
985 lz->func = func;
986 lz->it = it;
987 lz->stop = 0;
988
989 return (PyObject *)lz;
990}
991
992static void
993takewhile_dealloc(takewhileobject *lz)
994{
995 PyObject_GC_UnTrack(lz);
996 Py_XDECREF(lz->func);
997 Py_XDECREF(lz->it);
998 lz->ob_type->tp_free(lz);
999}
1000
1001static int
1002takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1003{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001004 Py_VISIT(lz->it);
1005 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001006 return 0;
1007}
1008
1009static PyObject *
1010takewhile_next(takewhileobject *lz)
1011{
1012 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001013 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001014 long ok;
1015
1016 if (lz->stop == 1)
1017 return NULL;
1018
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001019 assert(PyIter_Check(it));
1020 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001021 if (item == NULL)
1022 return NULL;
1023
1024 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1025 if (good == NULL) {
1026 Py_DECREF(item);
1027 return NULL;
1028 }
1029 ok = PyObject_IsTrue(good);
1030 Py_DECREF(good);
1031 if (ok)
1032 return item;
1033 Py_DECREF(item);
1034 lz->stop = 1;
1035 return NULL;
1036}
1037
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001038PyDoc_STRVAR(takewhile_doc,
1039"takewhile(predicate, iterable) --> takewhile object\n\
1040\n\
1041Return successive entries from an iterable as long as the \n\
1042predicate evaluates to true for each entry.");
1043
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001044static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001045 PyObject_HEAD_INIT(NULL)
1046 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001047 "itertools.takewhile", /* tp_name */
1048 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001049 0, /* tp_itemsize */
1050 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001051 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001052 0, /* tp_print */
1053 0, /* tp_getattr */
1054 0, /* tp_setattr */
1055 0, /* tp_compare */
1056 0, /* tp_repr */
1057 0, /* tp_as_number */
1058 0, /* tp_as_sequence */
1059 0, /* tp_as_mapping */
1060 0, /* tp_hash */
1061 0, /* tp_call */
1062 0, /* tp_str */
1063 PyObject_GenericGetAttr, /* tp_getattro */
1064 0, /* tp_setattro */
1065 0, /* tp_as_buffer */
1066 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1067 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001068 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001069 (traverseproc)takewhile_traverse, /* tp_traverse */
1070 0, /* tp_clear */
1071 0, /* tp_richcompare */
1072 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001073 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001074 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075 0, /* tp_methods */
1076 0, /* tp_members */
1077 0, /* tp_getset */
1078 0, /* tp_base */
1079 0, /* tp_dict */
1080 0, /* tp_descr_get */
1081 0, /* tp_descr_set */
1082 0, /* tp_dictoffset */
1083 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001084 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001085 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001086 PyObject_GC_Del, /* tp_free */
1087};
1088
1089
1090/* islice object ************************************************************/
1091
1092typedef struct {
1093 PyObject_HEAD
1094 PyObject *it;
1095 long next;
1096 long stop;
1097 long step;
1098 long cnt;
1099} isliceobject;
1100
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001101static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001102
1103static PyObject *
1104islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1105{
1106 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001107 long start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001108 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001109 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001110 isliceobject *lz;
1111
Georg Brandl02c42872005-08-26 06:42:30 +00001112 if (!_PyArg_NoKeywords("islice()", kwds))
1113 return NULL;
1114
Raymond Hettingerb2594052004-12-05 09:25:51 +00001115 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001116 return NULL;
1117
Raymond Hettingerb2594052004-12-05 09:25:51 +00001118 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001119 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001120 if (a1 != Py_None) {
1121 stop = PyInt_AsLong(a1);
1122 if (stop == -1) {
1123 if (PyErr_Occurred())
1124 PyErr_Clear();
1125 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001126 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001127 return NULL;
1128 }
1129 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001130 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001131 if (a1 != Py_None)
1132 start = PyInt_AsLong(a1);
1133 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001134 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001135 if (a2 != Py_None) {
1136 stop = PyInt_AsLong(a2);
1137 if (stop == -1) {
1138 if (PyErr_Occurred())
1139 PyErr_Clear();
1140 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001141 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001142 return NULL;
1143 }
1144 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001145 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001146 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001147 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001148 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001149 return NULL;
1150 }
1151
Raymond Hettingerb2594052004-12-05 09:25:51 +00001152 if (a3 != NULL) {
1153 if (a3 != Py_None)
1154 step = PyInt_AsLong(a3);
1155 if (step == -1 && PyErr_Occurred())
1156 PyErr_Clear();
1157 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001158 if (step<1) {
1159 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001160 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161 return NULL;
1162 }
1163
1164 /* Get iterator. */
1165 it = PyObject_GetIter(seq);
1166 if (it == NULL)
1167 return NULL;
1168
1169 /* create isliceobject structure */
1170 lz = (isliceobject *)type->tp_alloc(type, 0);
1171 if (lz == NULL) {
1172 Py_DECREF(it);
1173 return NULL;
1174 }
1175 lz->it = it;
1176 lz->next = start;
1177 lz->stop = stop;
1178 lz->step = step;
1179 lz->cnt = 0L;
1180
1181 return (PyObject *)lz;
1182}
1183
1184static void
1185islice_dealloc(isliceobject *lz)
1186{
1187 PyObject_GC_UnTrack(lz);
1188 Py_XDECREF(lz->it);
1189 lz->ob_type->tp_free(lz);
1190}
1191
1192static int
1193islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1194{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001195 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196 return 0;
1197}
1198
1199static PyObject *
1200islice_next(isliceobject *lz)
1201{
1202 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001203 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001204 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001205 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001206
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001207 assert(PyIter_Check(it));
1208 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001209 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001210 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001211 if (item == NULL)
1212 return NULL;
1213 Py_DECREF(item);
1214 lz->cnt++;
1215 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001216 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001218 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001219 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220 if (item == NULL)
1221 return NULL;
1222 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001223 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001225 if (lz->next < oldnext) /* Check for overflow */
1226 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001227 return item;
1228}
1229
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001230PyDoc_STRVAR(islice_doc,
1231"islice(iterable, [start,] stop [, step]) --> islice object\n\
1232\n\
1233Return an iterator whose next() method returns selected values from an\n\
1234iterable. If start is specified, will skip all preceding elements;\n\
1235otherwise, start defaults to zero. Step defaults to one. If\n\
1236specified as another value, step determines how many values are \n\
1237skipped between successive calls. Works like a slice() on a list\n\
1238but returns an iterator.");
1239
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001240static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001241 PyObject_HEAD_INIT(NULL)
1242 0, /* ob_size */
1243 "itertools.islice", /* tp_name */
1244 sizeof(isliceobject), /* tp_basicsize */
1245 0, /* tp_itemsize */
1246 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001247 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001248 0, /* tp_print */
1249 0, /* tp_getattr */
1250 0, /* tp_setattr */
1251 0, /* tp_compare */
1252 0, /* tp_repr */
1253 0, /* tp_as_number */
1254 0, /* tp_as_sequence */
1255 0, /* tp_as_mapping */
1256 0, /* tp_hash */
1257 0, /* tp_call */
1258 0, /* tp_str */
1259 PyObject_GenericGetAttr, /* tp_getattro */
1260 0, /* tp_setattro */
1261 0, /* tp_as_buffer */
1262 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1263 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001264 islice_doc, /* tp_doc */
1265 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001266 0, /* tp_clear */
1267 0, /* tp_richcompare */
1268 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001269 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001270 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001271 0, /* tp_methods */
1272 0, /* tp_members */
1273 0, /* tp_getset */
1274 0, /* tp_base */
1275 0, /* tp_dict */
1276 0, /* tp_descr_get */
1277 0, /* tp_descr_set */
1278 0, /* tp_dictoffset */
1279 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001280 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001281 islice_new, /* tp_new */
1282 PyObject_GC_Del, /* tp_free */
1283};
1284
1285
1286/* starmap object ************************************************************/
1287
1288typedef struct {
1289 PyObject_HEAD
1290 PyObject *func;
1291 PyObject *it;
1292} starmapobject;
1293
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001294static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001295
1296static PyObject *
1297starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1298{
1299 PyObject *func, *seq;
1300 PyObject *it;
1301 starmapobject *lz;
1302
Georg Brandl02c42872005-08-26 06:42:30 +00001303 if (!_PyArg_NoKeywords("starmap()", kwds))
1304 return NULL;
1305
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1307 return NULL;
1308
1309 /* Get iterator. */
1310 it = PyObject_GetIter(seq);
1311 if (it == NULL)
1312 return NULL;
1313
1314 /* create starmapobject structure */
1315 lz = (starmapobject *)type->tp_alloc(type, 0);
1316 if (lz == NULL) {
1317 Py_DECREF(it);
1318 return NULL;
1319 }
1320 Py_INCREF(func);
1321 lz->func = func;
1322 lz->it = it;
1323
1324 return (PyObject *)lz;
1325}
1326
1327static void
1328starmap_dealloc(starmapobject *lz)
1329{
1330 PyObject_GC_UnTrack(lz);
1331 Py_XDECREF(lz->func);
1332 Py_XDECREF(lz->it);
1333 lz->ob_type->tp_free(lz);
1334}
1335
1336static int
1337starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1338{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001339 Py_VISIT(lz->it);
1340 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341 return 0;
1342}
1343
1344static PyObject *
1345starmap_next(starmapobject *lz)
1346{
1347 PyObject *args;
1348 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001349 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001351 assert(PyIter_Check(it));
1352 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353 if (args == NULL)
1354 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001355 if (!PyTuple_CheckExact(args)) {
1356 Py_DECREF(args);
1357 PyErr_SetString(PyExc_TypeError,
1358 "iterator must return a tuple");
1359 return NULL;
1360 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361 result = PyObject_Call(lz->func, args, NULL);
1362 Py_DECREF(args);
1363 return result;
1364}
1365
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366PyDoc_STRVAR(starmap_doc,
1367"starmap(function, sequence) --> starmap object\n\
1368\n\
1369Return an iterator whose values are returned from the function evaluated\n\
1370with a argument tuple taken from the given sequence.");
1371
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001372static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001373 PyObject_HEAD_INIT(NULL)
1374 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001375 "itertools.starmap", /* tp_name */
1376 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377 0, /* tp_itemsize */
1378 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001379 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380 0, /* tp_print */
1381 0, /* tp_getattr */
1382 0, /* tp_setattr */
1383 0, /* tp_compare */
1384 0, /* tp_repr */
1385 0, /* tp_as_number */
1386 0, /* tp_as_sequence */
1387 0, /* tp_as_mapping */
1388 0, /* tp_hash */
1389 0, /* tp_call */
1390 0, /* tp_str */
1391 PyObject_GenericGetAttr, /* tp_getattro */
1392 0, /* tp_setattro */
1393 0, /* tp_as_buffer */
1394 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1395 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001396 starmap_doc, /* tp_doc */
1397 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398 0, /* tp_clear */
1399 0, /* tp_richcompare */
1400 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001401 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001402 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001403 0, /* tp_methods */
1404 0, /* tp_members */
1405 0, /* tp_getset */
1406 0, /* tp_base */
1407 0, /* tp_dict */
1408 0, /* tp_descr_get */
1409 0, /* tp_descr_set */
1410 0, /* tp_dictoffset */
1411 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001412 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001413 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001414 PyObject_GC_Del, /* tp_free */
1415};
1416
1417
1418/* imap object ************************************************************/
1419
1420typedef struct {
1421 PyObject_HEAD
1422 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423 PyObject *func;
1424} imapobject;
1425
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001426static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427
1428static PyObject *
1429imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1430{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001431 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001433 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001434
Georg Brandl02c42872005-08-26 06:42:30 +00001435 if (!_PyArg_NoKeywords("imap()", kwds))
1436 return NULL;
1437
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001438 numargs = PyTuple_Size(args);
1439 if (numargs < 2) {
1440 PyErr_SetString(PyExc_TypeError,
1441 "imap() must have at least two arguments.");
1442 return NULL;
1443 }
1444
1445 iters = PyTuple_New(numargs-1);
1446 if (iters == NULL)
1447 return NULL;
1448
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001449 for (i=1 ; i<numargs ; i++) {
1450 /* Get iterator. */
1451 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1452 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001453 Py_DECREF(iters);
1454 return NULL;
1455 }
1456 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457 }
1458
1459 /* create imapobject structure */
1460 lz = (imapobject *)type->tp_alloc(type, 0);
1461 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462 Py_DECREF(iters);
1463 return NULL;
1464 }
1465 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466 func = PyTuple_GET_ITEM(args, 0);
1467 Py_INCREF(func);
1468 lz->func = func;
1469
1470 return (PyObject *)lz;
1471}
1472
1473static void
1474imap_dealloc(imapobject *lz)
1475{
1476 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001477 Py_XDECREF(lz->iters);
1478 Py_XDECREF(lz->func);
1479 lz->ob_type->tp_free(lz);
1480}
1481
1482static int
1483imap_traverse(imapobject *lz, visitproc visit, void *arg)
1484{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001485 Py_VISIT(lz->iters);
1486 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487 return 0;
1488}
1489
Raymond Hettinger2012f172003-02-07 05:32:58 +00001490/*
1491imap() is an iterator version of __builtins__.map() except that it does
1492not have the None fill-in feature. That was intentionally left out for
1493the following reasons:
1494
1495 1) Itertools are designed to be easily combined and chained together.
1496 Having all tools stop with the shortest input is a unifying principle
1497 that makes it easier to combine finite iterators (supplying data) with
1498 infinite iterators like count() and repeat() (for supplying sequential
1499 or constant arguments to a function).
1500
1501 2) In typical use cases for combining itertools, having one finite data
1502 supplier run out before another is likely to be an error condition which
1503 should not pass silently by automatically supplying None.
1504
1505 3) The use cases for automatic None fill-in are rare -- not many functions
1506 do something useful when a parameter suddenly switches type and becomes
1507 None.
1508
1509 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001510 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001511
1512 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1513*/
1514
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001515static PyObject *
1516imap_next(imapobject *lz)
1517{
1518 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001519 PyObject *argtuple;
1520 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001521 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001522
1523 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001524 argtuple = PyTuple_New(numargs);
1525 if (argtuple == NULL)
1526 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001527
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001528 for (i=0 ; i<numargs ; i++) {
1529 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1530 if (val == NULL) {
1531 Py_DECREF(argtuple);
1532 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001534 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001536 if (lz->func == Py_None)
1537 return argtuple;
1538 result = PyObject_Call(lz->func, argtuple, NULL);
1539 Py_DECREF(argtuple);
1540 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541}
1542
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543PyDoc_STRVAR(imap_doc,
1544"imap(func, *iterables) --> imap object\n\
1545\n\
1546Make an iterator that computes the function using arguments from\n\
1547each of the iterables. Like map() except that it returns\n\
1548an iterator instead of a list and that it stops when the shortest\n\
1549iterable is exhausted instead of filling in None for shorter\n\
1550iterables.");
1551
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001552static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553 PyObject_HEAD_INIT(NULL)
1554 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001555 "itertools.imap", /* tp_name */
1556 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001557 0, /* tp_itemsize */
1558 /* methods */
1559 (destructor)imap_dealloc, /* tp_dealloc */
1560 0, /* tp_print */
1561 0, /* tp_getattr */
1562 0, /* tp_setattr */
1563 0, /* tp_compare */
1564 0, /* tp_repr */
1565 0, /* tp_as_number */
1566 0, /* tp_as_sequence */
1567 0, /* tp_as_mapping */
1568 0, /* tp_hash */
1569 0, /* tp_call */
1570 0, /* tp_str */
1571 PyObject_GenericGetAttr, /* tp_getattro */
1572 0, /* tp_setattro */
1573 0, /* tp_as_buffer */
1574 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1575 Py_TPFLAGS_BASETYPE, /* tp_flags */
1576 imap_doc, /* tp_doc */
1577 (traverseproc)imap_traverse, /* tp_traverse */
1578 0, /* tp_clear */
1579 0, /* tp_richcompare */
1580 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001581 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001582 (iternextfunc)imap_next, /* tp_iternext */
1583 0, /* tp_methods */
1584 0, /* tp_members */
1585 0, /* tp_getset */
1586 0, /* tp_base */
1587 0, /* tp_dict */
1588 0, /* tp_descr_get */
1589 0, /* tp_descr_set */
1590 0, /* tp_dictoffset */
1591 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001592 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001593 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001594 PyObject_GC_Del, /* tp_free */
1595};
1596
1597
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001598/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001599
1600typedef struct {
1601 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001602 Py_ssize_t tuplesize;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001603 long iternum; /* which iterator is active */
1604 PyObject *ittuple; /* tuple of iterators */
1605} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001606
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001607static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608
1609static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001610chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001612 chainobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001613 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001614 int i;
1615 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001616
Georg Brandl02c42872005-08-26 06:42:30 +00001617 if (!_PyArg_NoKeywords("chain()", kwds))
1618 return NULL;
1619
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001620 /* obtain iterators */
1621 assert(PyTuple_Check(args));
1622 ittuple = PyTuple_New(tuplesize);
1623 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001624 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001625 for (i=0; i < tuplesize; ++i) {
1626 PyObject *item = PyTuple_GET_ITEM(args, i);
1627 PyObject *it = PyObject_GetIter(item);
1628 if (it == NULL) {
1629 if (PyErr_ExceptionMatches(PyExc_TypeError))
1630 PyErr_Format(PyExc_TypeError,
1631 "chain argument #%d must support iteration",
1632 i+1);
1633 Py_DECREF(ittuple);
1634 return NULL;
1635 }
1636 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001637 }
1638
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001639 /* create chainobject structure */
1640 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001641 if (lz == NULL) {
1642 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001643 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001644 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001645
1646 lz->ittuple = ittuple;
1647 lz->iternum = 0;
1648 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001649
1650 return (PyObject *)lz;
1651}
1652
1653static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001654chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001655{
1656 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001657 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658 lz->ob_type->tp_free(lz);
1659}
1660
Raymond Hettinger2012f172003-02-07 05:32:58 +00001661static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001662chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001663{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001664 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001665 return 0;
1666}
1667
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001669chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001671 PyObject *it;
1672 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001673
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001674 while (lz->iternum < lz->tuplesize) {
1675 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1676 item = PyIter_Next(it);
1677 if (item != NULL)
1678 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001679 if (PyErr_Occurred()) {
1680 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1681 PyErr_Clear();
1682 else
1683 return NULL;
1684 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001685 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001686 }
1687 return NULL;
1688}
1689
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001690PyDoc_STRVAR(chain_doc,
1691"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001692\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001693Return a chain object whose .next() method returns elements from the\n\
1694first iterable until it is exhausted, then elements from the next\n\
1695iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001696
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001697static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001698 PyObject_HEAD_INIT(NULL)
1699 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001700 "itertools.chain", /* tp_name */
1701 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001702 0, /* tp_itemsize */
1703 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001704 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705 0, /* tp_print */
1706 0, /* tp_getattr */
1707 0, /* tp_setattr */
1708 0, /* tp_compare */
1709 0, /* tp_repr */
1710 0, /* tp_as_number */
1711 0, /* tp_as_sequence */
1712 0, /* tp_as_mapping */
1713 0, /* tp_hash */
1714 0, /* tp_call */
1715 0, /* tp_str */
1716 PyObject_GenericGetAttr, /* tp_getattro */
1717 0, /* tp_setattro */
1718 0, /* tp_as_buffer */
1719 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1720 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001721 chain_doc, /* tp_doc */
1722 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001723 0, /* tp_clear */
1724 0, /* tp_richcompare */
1725 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001726 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001727 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001728 0, /* tp_methods */
1729 0, /* tp_members */
1730 0, /* tp_getset */
1731 0, /* tp_base */
1732 0, /* tp_dict */
1733 0, /* tp_descr_get */
1734 0, /* tp_descr_set */
1735 0, /* tp_dictoffset */
1736 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001737 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001738 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739 PyObject_GC_Del, /* tp_free */
1740};
1741
1742
1743/* ifilter object ************************************************************/
1744
1745typedef struct {
1746 PyObject_HEAD
1747 PyObject *func;
1748 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749} ifilterobject;
1750
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001751static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001752
1753static PyObject *
1754ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1755{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001756 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757 PyObject *it;
1758 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759
Georg Brandl02c42872005-08-26 06:42:30 +00001760 if (!_PyArg_NoKeywords("ifilter()", kwds))
1761 return NULL;
1762
Raymond Hettinger60eca932003-02-09 06:40:58 +00001763 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764 return NULL;
1765
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001766 /* Get iterator. */
1767 it = PyObject_GetIter(seq);
1768 if (it == NULL)
1769 return NULL;
1770
1771 /* create ifilterobject structure */
1772 lz = (ifilterobject *)type->tp_alloc(type, 0);
1773 if (lz == NULL) {
1774 Py_DECREF(it);
1775 return NULL;
1776 }
1777 Py_INCREF(func);
1778 lz->func = func;
1779 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001780
1781 return (PyObject *)lz;
1782}
1783
1784static void
1785ifilter_dealloc(ifilterobject *lz)
1786{
1787 PyObject_GC_UnTrack(lz);
1788 Py_XDECREF(lz->func);
1789 Py_XDECREF(lz->it);
1790 lz->ob_type->tp_free(lz);
1791}
1792
1793static int
1794ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1795{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001796 Py_VISIT(lz->it);
1797 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001798 return 0;
1799}
1800
1801static PyObject *
1802ifilter_next(ifilterobject *lz)
1803{
1804 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001805 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001807 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001809 assert(PyIter_Check(it));
1810 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001811 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001812 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001813 if (item == NULL)
1814 return NULL;
1815
1816 if (lz->func == Py_None) {
1817 ok = PyObject_IsTrue(item);
1818 } else {
1819 PyObject *good;
1820 good = PyObject_CallFunctionObjArgs(lz->func,
1821 item, NULL);
1822 if (good == NULL) {
1823 Py_DECREF(item);
1824 return NULL;
1825 }
1826 ok = PyObject_IsTrue(good);
1827 Py_DECREF(good);
1828 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001829 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830 return item;
1831 Py_DECREF(item);
1832 }
1833}
1834
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001835PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001836"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001837\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001838Return those items of sequence for which function(item) is true.\n\
1839If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001840
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001841static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001842 PyObject_HEAD_INIT(NULL)
1843 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001844 "itertools.ifilter", /* tp_name */
1845 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001846 0, /* tp_itemsize */
1847 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001848 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001849 0, /* tp_print */
1850 0, /* tp_getattr */
1851 0, /* tp_setattr */
1852 0, /* tp_compare */
1853 0, /* tp_repr */
1854 0, /* tp_as_number */
1855 0, /* tp_as_sequence */
1856 0, /* tp_as_mapping */
1857 0, /* tp_hash */
1858 0, /* tp_call */
1859 0, /* tp_str */
1860 PyObject_GenericGetAttr, /* tp_getattro */
1861 0, /* tp_setattro */
1862 0, /* tp_as_buffer */
1863 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1864 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001865 ifilter_doc, /* tp_doc */
1866 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867 0, /* tp_clear */
1868 0, /* tp_richcompare */
1869 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001870 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001871 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001872 0, /* tp_methods */
1873 0, /* tp_members */
1874 0, /* tp_getset */
1875 0, /* tp_base */
1876 0, /* tp_dict */
1877 0, /* tp_descr_get */
1878 0, /* tp_descr_set */
1879 0, /* tp_dictoffset */
1880 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001881 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001882 ifilter_new, /* tp_new */
1883 PyObject_GC_Del, /* tp_free */
1884};
1885
1886
1887/* ifilterfalse object ************************************************************/
1888
1889typedef struct {
1890 PyObject_HEAD
1891 PyObject *func;
1892 PyObject *it;
1893} ifilterfalseobject;
1894
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001895static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001896
1897static PyObject *
1898ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1899{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001900 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001901 PyObject *it;
1902 ifilterfalseobject *lz;
1903
Georg Brandl02c42872005-08-26 06:42:30 +00001904 if (!_PyArg_NoKeywords("ifilterfalse()", kwds))
1905 return NULL;
1906
Raymond Hettinger60eca932003-02-09 06:40:58 +00001907 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1908 return NULL;
1909
1910 /* Get iterator. */
1911 it = PyObject_GetIter(seq);
1912 if (it == NULL)
1913 return NULL;
1914
1915 /* create ifilterfalseobject structure */
1916 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1917 if (lz == NULL) {
1918 Py_DECREF(it);
1919 return NULL;
1920 }
1921 Py_INCREF(func);
1922 lz->func = func;
1923 lz->it = it;
1924
1925 return (PyObject *)lz;
1926}
1927
1928static void
1929ifilterfalse_dealloc(ifilterfalseobject *lz)
1930{
1931 PyObject_GC_UnTrack(lz);
1932 Py_XDECREF(lz->func);
1933 Py_XDECREF(lz->it);
1934 lz->ob_type->tp_free(lz);
1935}
1936
1937static int
1938ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1939{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001940 Py_VISIT(lz->it);
1941 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001942 return 0;
1943}
1944
1945static PyObject *
1946ifilterfalse_next(ifilterfalseobject *lz)
1947{
1948 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001949 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001950 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001951 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001952
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001953 assert(PyIter_Check(it));
1954 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001955 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001956 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001957 if (item == NULL)
1958 return NULL;
1959
1960 if (lz->func == Py_None) {
1961 ok = PyObject_IsTrue(item);
1962 } else {
1963 PyObject *good;
1964 good = PyObject_CallFunctionObjArgs(lz->func,
1965 item, NULL);
1966 if (good == NULL) {
1967 Py_DECREF(item);
1968 return NULL;
1969 }
1970 ok = PyObject_IsTrue(good);
1971 Py_DECREF(good);
1972 }
1973 if (!ok)
1974 return item;
1975 Py_DECREF(item);
1976 }
1977}
1978
Raymond Hettinger60eca932003-02-09 06:40:58 +00001979PyDoc_STRVAR(ifilterfalse_doc,
1980"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1981\n\
1982Return those items of sequence for which function(item) is false.\n\
1983If function is None, return the items that are false.");
1984
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001985static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001986 PyObject_HEAD_INIT(NULL)
1987 0, /* ob_size */
1988 "itertools.ifilterfalse", /* tp_name */
1989 sizeof(ifilterfalseobject), /* tp_basicsize */
1990 0, /* tp_itemsize */
1991 /* methods */
1992 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1993 0, /* tp_print */
1994 0, /* tp_getattr */
1995 0, /* tp_setattr */
1996 0, /* tp_compare */
1997 0, /* tp_repr */
1998 0, /* tp_as_number */
1999 0, /* tp_as_sequence */
2000 0, /* tp_as_mapping */
2001 0, /* tp_hash */
2002 0, /* tp_call */
2003 0, /* tp_str */
2004 PyObject_GenericGetAttr, /* tp_getattro */
2005 0, /* tp_setattro */
2006 0, /* tp_as_buffer */
2007 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2008 Py_TPFLAGS_BASETYPE, /* tp_flags */
2009 ifilterfalse_doc, /* tp_doc */
2010 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2011 0, /* tp_clear */
2012 0, /* tp_richcompare */
2013 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002014 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002015 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2016 0, /* tp_methods */
2017 0, /* tp_members */
2018 0, /* tp_getset */
2019 0, /* tp_base */
2020 0, /* tp_dict */
2021 0, /* tp_descr_get */
2022 0, /* tp_descr_set */
2023 0, /* tp_dictoffset */
2024 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002025 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002026 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002027 PyObject_GC_Del, /* tp_free */
2028};
2029
2030
2031/* count object ************************************************************/
2032
2033typedef struct {
2034 PyObject_HEAD
2035 long cnt;
2036} countobject;
2037
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002038static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002039
2040static PyObject *
2041count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2042{
2043 countobject *lz;
2044 long cnt = 0;
2045
Georg Brandl02c42872005-08-26 06:42:30 +00002046 if (!_PyArg_NoKeywords("count()", kwds))
2047 return NULL;
2048
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002049 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
2050 return NULL;
2051
2052 /* create countobject structure */
2053 lz = (countobject *)PyObject_New(countobject, &count_type);
2054 if (lz == NULL)
2055 return NULL;
2056 lz->cnt = cnt;
2057
2058 return (PyObject *)lz;
2059}
2060
2061static PyObject *
2062count_next(countobject *lz)
2063{
2064 return PyInt_FromLong(lz->cnt++);
2065}
2066
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002067static PyObject *
2068count_repr(countobject *lz)
2069{
Brett Cannon00468392004-04-13 02:43:53 +00002070 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002071}
2072
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002073PyDoc_STRVAR(count_doc,
2074"count([firstval]) --> count object\n\
2075\n\
2076Return a count object whose .next() method returns consecutive\n\
2077integers starting from zero or, if specified, from firstval.");
2078
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002079static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002080 PyObject_HEAD_INIT(NULL)
2081 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002082 "itertools.count", /* tp_name */
2083 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002084 0, /* tp_itemsize */
2085 /* methods */
2086 (destructor)PyObject_Del, /* tp_dealloc */
2087 0, /* tp_print */
2088 0, /* tp_getattr */
2089 0, /* tp_setattr */
2090 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002091 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002092 0, /* tp_as_number */
2093 0, /* tp_as_sequence */
2094 0, /* tp_as_mapping */
2095 0, /* tp_hash */
2096 0, /* tp_call */
2097 0, /* tp_str */
2098 PyObject_GenericGetAttr, /* tp_getattro */
2099 0, /* tp_setattro */
2100 0, /* tp_as_buffer */
2101 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002102 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002103 0, /* tp_traverse */
2104 0, /* tp_clear */
2105 0, /* tp_richcompare */
2106 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002107 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002108 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002109 0, /* tp_methods */
2110 0, /* tp_members */
2111 0, /* tp_getset */
2112 0, /* tp_base */
2113 0, /* tp_dict */
2114 0, /* tp_descr_get */
2115 0, /* tp_descr_set */
2116 0, /* tp_dictoffset */
2117 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002118 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002119 count_new, /* tp_new */
2120};
2121
2122
2123/* izip object ************************************************************/
2124
2125#include "Python.h"
2126
2127typedef struct {
2128 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002129 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002130 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002131 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002132} izipobject;
2133
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002134static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002135
2136static PyObject *
2137izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2138{
2139 izipobject *lz;
2140 int i;
2141 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002142 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002143 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002144
Georg Brandl02c42872005-08-26 06:42:30 +00002145 if (!_PyArg_NoKeywords("izip()", kwds))
2146 return NULL;
2147
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002148 /* args must be a tuple */
2149 assert(PyTuple_Check(args));
2150
2151 /* obtain iterators */
2152 ittuple = PyTuple_New(tuplesize);
2153 if(ittuple == NULL)
2154 return NULL;
2155 for (i=0; i < tuplesize; ++i) {
2156 PyObject *item = PyTuple_GET_ITEM(args, i);
2157 PyObject *it = PyObject_GetIter(item);
2158 if (it == NULL) {
2159 if (PyErr_ExceptionMatches(PyExc_TypeError))
2160 PyErr_Format(PyExc_TypeError,
2161 "izip argument #%d must support iteration",
2162 i+1);
2163 Py_DECREF(ittuple);
2164 return NULL;
2165 }
2166 PyTuple_SET_ITEM(ittuple, i, it);
2167 }
2168
Raymond Hettinger2012f172003-02-07 05:32:58 +00002169 /* create a result holder */
2170 result = PyTuple_New(tuplesize);
2171 if (result == NULL) {
2172 Py_DECREF(ittuple);
2173 return NULL;
2174 }
2175 for (i=0 ; i < tuplesize ; i++) {
2176 Py_INCREF(Py_None);
2177 PyTuple_SET_ITEM(result, i, Py_None);
2178 }
2179
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002180 /* create izipobject structure */
2181 lz = (izipobject *)type->tp_alloc(type, 0);
2182 if (lz == NULL) {
2183 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002184 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002185 return NULL;
2186 }
2187 lz->ittuple = ittuple;
2188 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002189 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002190
2191 return (PyObject *)lz;
2192}
2193
2194static void
2195izip_dealloc(izipobject *lz)
2196{
2197 PyObject_GC_UnTrack(lz);
2198 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002199 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002200 lz->ob_type->tp_free(lz);
2201}
2202
2203static int
2204izip_traverse(izipobject *lz, visitproc visit, void *arg)
2205{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002206 Py_VISIT(lz->ittuple);
2207 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002208 return 0;
2209}
2210
2211static PyObject *
2212izip_next(izipobject *lz)
2213{
2214 int i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002215 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002216 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002217 PyObject *it;
2218 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002219 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002220
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002221 if (tuplesize == 0)
2222 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002223 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002224 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002225 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002226 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002227 assert(PyIter_Check(it));
2228 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002229 if (item == NULL) {
2230 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002231 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002232 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002233 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002234 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002235 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002236 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002237 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002238 result = PyTuple_New(tuplesize);
2239 if (result == NULL)
2240 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002241 for (i=0 ; i < tuplesize ; i++) {
2242 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002243 assert(PyIter_Check(it));
2244 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002245 if (item == NULL) {
2246 Py_DECREF(result);
2247 return NULL;
2248 }
2249 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002250 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002251 }
2252 return result;
2253}
2254
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002255PyDoc_STRVAR(izip_doc,
2256"izip(iter1 [,iter2 [...]]) --> izip object\n\
2257\n\
2258Return a izip object whose .next() method returns a tuple where\n\
2259the i-th element comes from the i-th iterable argument. The .next()\n\
2260method continues until the shortest iterable in the argument sequence\n\
2261is exhausted and then it raises StopIteration. Works like the zip()\n\
2262function but consumes less memory by returning an iterator instead of\n\
2263a list.");
2264
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002265static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002266 PyObject_HEAD_INIT(NULL)
2267 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002268 "itertools.izip", /* tp_name */
2269 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002270 0, /* tp_itemsize */
2271 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002272 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002273 0, /* tp_print */
2274 0, /* tp_getattr */
2275 0, /* tp_setattr */
2276 0, /* tp_compare */
2277 0, /* tp_repr */
2278 0, /* tp_as_number */
2279 0, /* tp_as_sequence */
2280 0, /* tp_as_mapping */
2281 0, /* tp_hash */
2282 0, /* tp_call */
2283 0, /* tp_str */
2284 PyObject_GenericGetAttr, /* tp_getattro */
2285 0, /* tp_setattro */
2286 0, /* tp_as_buffer */
2287 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2288 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002289 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002290 (traverseproc)izip_traverse, /* tp_traverse */
2291 0, /* tp_clear */
2292 0, /* tp_richcompare */
2293 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002294 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002295 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002296 0, /* tp_methods */
2297 0, /* tp_members */
2298 0, /* tp_getset */
2299 0, /* tp_base */
2300 0, /* tp_dict */
2301 0, /* tp_descr_get */
2302 0, /* tp_descr_set */
2303 0, /* tp_dictoffset */
2304 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002305 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002306 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002307 PyObject_GC_Del, /* tp_free */
2308};
2309
2310
2311/* repeat object ************************************************************/
2312
2313typedef struct {
2314 PyObject_HEAD
2315 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002316 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002317} repeatobject;
2318
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002319static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002320
2321static PyObject *
2322repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2323{
2324 repeatobject *ro;
2325 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002326 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002327
Georg Brandl02c42872005-08-26 06:42:30 +00002328 if (!_PyArg_NoKeywords("repeat()", kwds))
2329 return NULL;
2330
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002331 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002332 return NULL;
2333
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002334 if (PyTuple_Size(args) == 2 && cnt < 0)
2335 cnt = 0;
2336
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002337 ro = (repeatobject *)type->tp_alloc(type, 0);
2338 if (ro == NULL)
2339 return NULL;
2340 Py_INCREF(element);
2341 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002342 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002343 return (PyObject *)ro;
2344}
2345
2346static void
2347repeat_dealloc(repeatobject *ro)
2348{
2349 PyObject_GC_UnTrack(ro);
2350 Py_XDECREF(ro->element);
2351 ro->ob_type->tp_free(ro);
2352}
2353
2354static int
2355repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2356{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002357 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002358 return 0;
2359}
2360
2361static PyObject *
2362repeat_next(repeatobject *ro)
2363{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002364 if (ro->cnt == 0)
2365 return NULL;
2366 if (ro->cnt > 0)
2367 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002368 Py_INCREF(ro->element);
2369 return ro->element;
2370}
2371
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002372static PyObject *
2373repeat_repr(repeatobject *ro)
2374{
2375 PyObject *result, *objrepr;
2376
2377 objrepr = PyObject_Repr(ro->element);
2378 if (objrepr == NULL)
2379 return NULL;
2380
2381 if (ro->cnt == -1)
2382 result = PyString_FromFormat("repeat(%s)",
2383 PyString_AS_STRING(objrepr));
2384 else
Brett Cannon00468392004-04-13 02:43:53 +00002385 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002386 PyString_AS_STRING(objrepr), ro->cnt);
2387 Py_DECREF(objrepr);
2388 return result;
2389}
2390
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002391static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002392repeat_len(repeatobject *ro)
2393{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002394 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002395 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002396 return NULL;
2397 }
2398 return PyInt_FromLong(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002399}
2400
Armin Rigof5b3e362006-02-11 21:32:43 +00002401PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002402
2403static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002404 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002405 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002406};
2407
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002408PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002409"repeat(element [,times]) -> create an iterator which returns the element\n\
2410for the specified number of times. If not specified, returns the element\n\
2411endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002412
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002413static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002414 PyObject_HEAD_INIT(NULL)
2415 0, /* ob_size */
2416 "itertools.repeat", /* tp_name */
2417 sizeof(repeatobject), /* tp_basicsize */
2418 0, /* tp_itemsize */
2419 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002420 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002421 0, /* tp_print */
2422 0, /* tp_getattr */
2423 0, /* tp_setattr */
2424 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002425 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002426 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002427 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002428 0, /* tp_as_mapping */
2429 0, /* tp_hash */
2430 0, /* tp_call */
2431 0, /* tp_str */
2432 PyObject_GenericGetAttr, /* tp_getattro */
2433 0, /* tp_setattro */
2434 0, /* tp_as_buffer */
2435 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2436 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002437 repeat_doc, /* tp_doc */
2438 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002439 0, /* tp_clear */
2440 0, /* tp_richcompare */
2441 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002442 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002443 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002444 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002445 0, /* tp_members */
2446 0, /* tp_getset */
2447 0, /* tp_base */
2448 0, /* tp_dict */
2449 0, /* tp_descr_get */
2450 0, /* tp_descr_set */
2451 0, /* tp_dictoffset */
2452 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002453 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002454 repeat_new, /* tp_new */
2455 PyObject_GC_Del, /* tp_free */
2456};
2457
2458
2459/* module level code ********************************************************/
2460
2461PyDoc_STRVAR(module_doc,
2462"Functional tools for creating and using iterators.\n\
2463\n\
2464Infinite iterators:\n\
2465count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002466cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002467repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002468\n\
2469Iterators terminating on the shortest input sequence:\n\
2470izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002471ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2472ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002473islice(seq, [start,] stop [, step]) --> elements from\n\
2474 seq[start:stop:step]\n\
2475imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2476starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002477tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002478chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002479takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2480dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002481groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002482");
2483
2484
Raymond Hettingerad983e72003-11-12 14:32:26 +00002485static PyMethodDef module_methods[] = {
2486 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2487 {NULL, NULL} /* sentinel */
2488};
2489
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002490PyMODINIT_FUNC
2491inititertools(void)
2492{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002493 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002494 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002495 char *name;
2496 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002497 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002498 &dropwhile_type,
2499 &takewhile_type,
2500 &islice_type,
2501 &starmap_type,
2502 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002503 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002504 &ifilter_type,
2505 &ifilterfalse_type,
2506 &count_type,
2507 &izip_type,
2508 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002509 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002510 NULL
2511 };
2512
Skip Montanarof3938fd2004-02-10 20:27:40 +00002513 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002514 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002515 if (m == NULL)
2516 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002517
Raymond Hettinger60eca932003-02-09 06:40:58 +00002518 for (i=0 ; typelist[i] != NULL ; i++) {
2519 if (PyType_Ready(typelist[i]) < 0)
2520 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002521 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002522 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002523 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002524 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002525 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002526
2527 if (PyType_Ready(&teedataobject_type) < 0)
2528 return;
2529 if (PyType_Ready(&tee_type) < 0)
2530 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002531 if (PyType_Ready(&_grouper_type) < 0)
2532 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002533}