blob: 86b1bbfb3404bd83db5e4d59ba35580017de59c0 [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;
Thomas Woutersb3deb942006-04-15 22:33:13 +0000501 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000502 return (PyObject *)newto;
503}
504
505PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
506
507static PyObject *
508tee_fromiterable(PyObject *iterable)
509{
510 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000511 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000512
513 it = PyObject_GetIter(iterable);
514 if (it == NULL)
515 return NULL;
516 if (PyObject_TypeCheck(it, &tee_type)) {
517 to = (teeobject *)tee_copy((teeobject *)it);
518 goto done;
519 }
520
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000521 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000522 if (to == NULL)
523 goto done;
524 to->dataobj = (teedataobject *)teedataobject_new(it);
525 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000526 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000527 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000528done:
529 Py_XDECREF(it);
530 return (PyObject *)to;
531}
532
533static PyObject *
534tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
535{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000536 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000537
538 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
539 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000540 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000541}
542
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000543static int
544tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000545{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000546 if (to->weakreflist != NULL)
547 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000548 Py_CLEAR(to->dataobj);
549 return 0;
550}
551
552static void
553tee_dealloc(teeobject *to)
554{
555 PyObject_GC_UnTrack(to);
556 tee_clear(to);
557 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000558}
559
Raymond Hettingerad983e72003-11-12 14:32:26 +0000560PyDoc_STRVAR(teeobject_doc,
561"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000562
Raymond Hettingerad983e72003-11-12 14:32:26 +0000563static PyMethodDef tee_methods[] = {
564 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
565 {NULL, NULL} /* sentinel */
566};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000567
568static PyTypeObject tee_type = {
569 PyObject_HEAD_INIT(NULL)
570 0, /* ob_size */
571 "itertools.tee", /* tp_name */
572 sizeof(teeobject), /* tp_basicsize */
573 0, /* tp_itemsize */
574 /* methods */
575 (destructor)tee_dealloc, /* tp_dealloc */
576 0, /* tp_print */
577 0, /* tp_getattr */
578 0, /* tp_setattr */
579 0, /* tp_compare */
580 0, /* tp_repr */
581 0, /* tp_as_number */
582 0, /* tp_as_sequence */
583 0, /* tp_as_mapping */
584 0, /* tp_hash */
585 0, /* tp_call */
586 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000587 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000588 0, /* tp_setattro */
589 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000590 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000591 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000592 (traverseproc)tee_traverse, /* tp_traverse */
593 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000594 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000595 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000596 PyObject_SelfIter, /* tp_iter */
597 (iternextfunc)tee_next, /* tp_iternext */
598 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000599 0, /* tp_members */
600 0, /* tp_getset */
601 0, /* tp_base */
602 0, /* tp_dict */
603 0, /* tp_descr_get */
604 0, /* tp_descr_set */
605 0, /* tp_dictoffset */
606 0, /* tp_init */
607 0, /* tp_alloc */
608 tee_new, /* tp_new */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000609 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000610};
611
Raymond Hettingerad983e72003-11-12 14:32:26 +0000612static PyObject *
613tee(PyObject *self, PyObject *args)
614{
615 int i, n=2;
616 PyObject *it, *iterable, *copyable, *result;
617
618 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
619 return NULL;
620 result = PyTuple_New(n);
621 if (result == NULL)
622 return NULL;
623 if (n == 0)
624 return result;
625 it = PyObject_GetIter(iterable);
626 if (it == NULL) {
627 Py_DECREF(result);
628 return NULL;
629 }
630 if (!PyObject_HasAttrString(it, "__copy__")) {
631 copyable = tee_fromiterable(it);
632 Py_DECREF(it);
633 if (copyable == NULL) {
634 Py_DECREF(result);
635 return NULL;
636 }
637 } else
638 copyable = it;
639 PyTuple_SET_ITEM(result, 0, copyable);
640 for (i=1 ; i<n ; i++) {
641 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
642 if (copyable == NULL) {
643 Py_DECREF(result);
644 return NULL;
645 }
646 PyTuple_SET_ITEM(result, i, copyable);
647 }
648 return result;
649}
650
651PyDoc_STRVAR(tee_doc,
652"tee(iterable, n=2) --> tuple of n independent iterators.");
653
654
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000655/* cycle object **********************************************************/
656
657typedef struct {
658 PyObject_HEAD
659 PyObject *it;
660 PyObject *saved;
661 int firstpass;
662} cycleobject;
663
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000664static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000665
666static PyObject *
667cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
668{
669 PyObject *it;
670 PyObject *iterable;
671 PyObject *saved;
672 cycleobject *lz;
673
Georg Brandl02c42872005-08-26 06:42:30 +0000674 if (!_PyArg_NoKeywords("cycle()", kwds))
675 return NULL;
676
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000677 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
678 return NULL;
679
680 /* Get iterator. */
681 it = PyObject_GetIter(iterable);
682 if (it == NULL)
683 return NULL;
684
685 saved = PyList_New(0);
686 if (saved == NULL) {
687 Py_DECREF(it);
688 return NULL;
689 }
690
691 /* create cycleobject structure */
692 lz = (cycleobject *)type->tp_alloc(type, 0);
693 if (lz == NULL) {
694 Py_DECREF(it);
695 Py_DECREF(saved);
696 return NULL;
697 }
698 lz->it = it;
699 lz->saved = saved;
700 lz->firstpass = 0;
701
702 return (PyObject *)lz;
703}
704
705static void
706cycle_dealloc(cycleobject *lz)
707{
708 PyObject_GC_UnTrack(lz);
709 Py_XDECREF(lz->saved);
710 Py_XDECREF(lz->it);
711 lz->ob_type->tp_free(lz);
712}
713
714static int
715cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
716{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000717 Py_VISIT(lz->it);
718 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000719 return 0;
720}
721
722static PyObject *
723cycle_next(cycleobject *lz)
724{
725 PyObject *item;
726 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000727 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000728
729 while (1) {
730 item = PyIter_Next(lz->it);
731 if (item != NULL) {
732 if (!lz->firstpass)
733 PyList_Append(lz->saved, item);
734 return item;
735 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000736 if (PyErr_Occurred()) {
737 if (PyErr_ExceptionMatches(PyExc_StopIteration))
738 PyErr_Clear();
739 else
740 return NULL;
741 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000742 if (PyList_Size(lz->saved) == 0)
743 return NULL;
744 it = PyObject_GetIter(lz->saved);
745 if (it == NULL)
746 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000747 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000748 lz->it = it;
749 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000750 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000751 }
752}
753
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000754PyDoc_STRVAR(cycle_doc,
755"cycle(iterable) --> cycle object\n\
756\n\
757Return elements from the iterable until it is exhausted.\n\
758Then repeat the sequence indefinitely.");
759
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000760static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000761 PyObject_HEAD_INIT(NULL)
762 0, /* ob_size */
763 "itertools.cycle", /* tp_name */
764 sizeof(cycleobject), /* tp_basicsize */
765 0, /* tp_itemsize */
766 /* methods */
767 (destructor)cycle_dealloc, /* tp_dealloc */
768 0, /* tp_print */
769 0, /* tp_getattr */
770 0, /* tp_setattr */
771 0, /* tp_compare */
772 0, /* tp_repr */
773 0, /* tp_as_number */
774 0, /* tp_as_sequence */
775 0, /* tp_as_mapping */
776 0, /* tp_hash */
777 0, /* tp_call */
778 0, /* tp_str */
779 PyObject_GenericGetAttr, /* tp_getattro */
780 0, /* tp_setattro */
781 0, /* tp_as_buffer */
782 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
783 Py_TPFLAGS_BASETYPE, /* tp_flags */
784 cycle_doc, /* tp_doc */
785 (traverseproc)cycle_traverse, /* tp_traverse */
786 0, /* tp_clear */
787 0, /* tp_richcompare */
788 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000789 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000790 (iternextfunc)cycle_next, /* tp_iternext */
791 0, /* tp_methods */
792 0, /* tp_members */
793 0, /* tp_getset */
794 0, /* tp_base */
795 0, /* tp_dict */
796 0, /* tp_descr_get */
797 0, /* tp_descr_set */
798 0, /* tp_dictoffset */
799 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000800 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000801 cycle_new, /* tp_new */
802 PyObject_GC_Del, /* tp_free */
803};
804
805
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000806/* dropwhile object **********************************************************/
807
808typedef struct {
809 PyObject_HEAD
810 PyObject *func;
811 PyObject *it;
812 long start;
813} dropwhileobject;
814
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000815static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000816
817static PyObject *
818dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
819{
820 PyObject *func, *seq;
821 PyObject *it;
822 dropwhileobject *lz;
823
Georg Brandl02c42872005-08-26 06:42:30 +0000824 if (!_PyArg_NoKeywords("dropwhile()", kwds))
825 return NULL;
826
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000827 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
828 return NULL;
829
830 /* Get iterator. */
831 it = PyObject_GetIter(seq);
832 if (it == NULL)
833 return NULL;
834
835 /* create dropwhileobject structure */
836 lz = (dropwhileobject *)type->tp_alloc(type, 0);
837 if (lz == NULL) {
838 Py_DECREF(it);
839 return NULL;
840 }
841 Py_INCREF(func);
842 lz->func = func;
843 lz->it = it;
844 lz->start = 0;
845
846 return (PyObject *)lz;
847}
848
849static void
850dropwhile_dealloc(dropwhileobject *lz)
851{
852 PyObject_GC_UnTrack(lz);
853 Py_XDECREF(lz->func);
854 Py_XDECREF(lz->it);
855 lz->ob_type->tp_free(lz);
856}
857
858static int
859dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
860{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000861 Py_VISIT(lz->it);
862 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000863 return 0;
864}
865
866static PyObject *
867dropwhile_next(dropwhileobject *lz)
868{
869 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000870 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000871 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000872 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000874 assert(PyIter_Check(it));
875 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000877 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878 if (item == NULL)
879 return NULL;
880 if (lz->start == 1)
881 return item;
882
883 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
884 if (good == NULL) {
885 Py_DECREF(item);
886 return NULL;
887 }
888 ok = PyObject_IsTrue(good);
889 Py_DECREF(good);
890 if (!ok) {
891 lz->start = 1;
892 return item;
893 }
894 Py_DECREF(item);
895 }
896}
897
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000898PyDoc_STRVAR(dropwhile_doc,
899"dropwhile(predicate, iterable) --> dropwhile object\n\
900\n\
901Drop items from the iterable while predicate(item) is true.\n\
902Afterwards, return every element until the iterable is exhausted.");
903
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000904static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000905 PyObject_HEAD_INIT(NULL)
906 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000907 "itertools.dropwhile", /* tp_name */
908 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000909 0, /* tp_itemsize */
910 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000911 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000912 0, /* tp_print */
913 0, /* tp_getattr */
914 0, /* tp_setattr */
915 0, /* tp_compare */
916 0, /* tp_repr */
917 0, /* tp_as_number */
918 0, /* tp_as_sequence */
919 0, /* tp_as_mapping */
920 0, /* tp_hash */
921 0, /* tp_call */
922 0, /* tp_str */
923 PyObject_GenericGetAttr, /* tp_getattro */
924 0, /* tp_setattro */
925 0, /* tp_as_buffer */
926 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
927 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000928 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000929 (traverseproc)dropwhile_traverse, /* tp_traverse */
930 0, /* tp_clear */
931 0, /* tp_richcompare */
932 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000933 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000934 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000935 0, /* tp_methods */
936 0, /* tp_members */
937 0, /* tp_getset */
938 0, /* tp_base */
939 0, /* tp_dict */
940 0, /* tp_descr_get */
941 0, /* tp_descr_set */
942 0, /* tp_dictoffset */
943 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000944 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000945 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000946 PyObject_GC_Del, /* tp_free */
947};
948
949
950/* takewhile object **********************************************************/
951
952typedef struct {
953 PyObject_HEAD
954 PyObject *func;
955 PyObject *it;
956 long stop;
957} takewhileobject;
958
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000959static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000960
961static PyObject *
962takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
963{
964 PyObject *func, *seq;
965 PyObject *it;
966 takewhileobject *lz;
967
Georg Brandl02c42872005-08-26 06:42:30 +0000968 if (!_PyArg_NoKeywords("takewhile()", kwds))
969 return NULL;
970
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000971 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
972 return NULL;
973
974 /* Get iterator. */
975 it = PyObject_GetIter(seq);
976 if (it == NULL)
977 return NULL;
978
979 /* create takewhileobject structure */
980 lz = (takewhileobject *)type->tp_alloc(type, 0);
981 if (lz == NULL) {
982 Py_DECREF(it);
983 return NULL;
984 }
985 Py_INCREF(func);
986 lz->func = func;
987 lz->it = it;
988 lz->stop = 0;
989
990 return (PyObject *)lz;
991}
992
993static void
994takewhile_dealloc(takewhileobject *lz)
995{
996 PyObject_GC_UnTrack(lz);
997 Py_XDECREF(lz->func);
998 Py_XDECREF(lz->it);
999 lz->ob_type->tp_free(lz);
1000}
1001
1002static int
1003takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1004{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001005 Py_VISIT(lz->it);
1006 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001007 return 0;
1008}
1009
1010static PyObject *
1011takewhile_next(takewhileobject *lz)
1012{
1013 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001014 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001015 long ok;
1016
1017 if (lz->stop == 1)
1018 return NULL;
1019
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001020 assert(PyIter_Check(it));
1021 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001022 if (item == NULL)
1023 return NULL;
1024
1025 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1026 if (good == NULL) {
1027 Py_DECREF(item);
1028 return NULL;
1029 }
1030 ok = PyObject_IsTrue(good);
1031 Py_DECREF(good);
1032 if (ok)
1033 return item;
1034 Py_DECREF(item);
1035 lz->stop = 1;
1036 return NULL;
1037}
1038
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001039PyDoc_STRVAR(takewhile_doc,
1040"takewhile(predicate, iterable) --> takewhile object\n\
1041\n\
1042Return successive entries from an iterable as long as the \n\
1043predicate evaluates to true for each entry.");
1044
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001045static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001046 PyObject_HEAD_INIT(NULL)
1047 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001048 "itertools.takewhile", /* tp_name */
1049 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001050 0, /* tp_itemsize */
1051 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001052 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053 0, /* tp_print */
1054 0, /* tp_getattr */
1055 0, /* tp_setattr */
1056 0, /* tp_compare */
1057 0, /* tp_repr */
1058 0, /* tp_as_number */
1059 0, /* tp_as_sequence */
1060 0, /* tp_as_mapping */
1061 0, /* tp_hash */
1062 0, /* tp_call */
1063 0, /* tp_str */
1064 PyObject_GenericGetAttr, /* tp_getattro */
1065 0, /* tp_setattro */
1066 0, /* tp_as_buffer */
1067 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1068 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001069 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001070 (traverseproc)takewhile_traverse, /* tp_traverse */
1071 0, /* tp_clear */
1072 0, /* tp_richcompare */
1073 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001074 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001075 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001076 0, /* tp_methods */
1077 0, /* tp_members */
1078 0, /* tp_getset */
1079 0, /* tp_base */
1080 0, /* tp_dict */
1081 0, /* tp_descr_get */
1082 0, /* tp_descr_set */
1083 0, /* tp_dictoffset */
1084 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001085 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001086 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001087 PyObject_GC_Del, /* tp_free */
1088};
1089
1090
1091/* islice object ************************************************************/
1092
1093typedef struct {
1094 PyObject_HEAD
1095 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001096 Py_ssize_t next;
1097 Py_ssize_t stop;
1098 Py_ssize_t step;
1099 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001100} isliceobject;
1101
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001102static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001103
1104static PyObject *
1105islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1106{
1107 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001108 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001109 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001110 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001111 isliceobject *lz;
1112
Georg Brandl02c42872005-08-26 06:42:30 +00001113 if (!_PyArg_NoKeywords("islice()", kwds))
1114 return NULL;
1115
Raymond Hettingerb2594052004-12-05 09:25:51 +00001116 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001117 return NULL;
1118
Raymond Hettingerb2594052004-12-05 09:25:51 +00001119 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001120 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001121 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001122 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001123 if (stop == -1) {
1124 if (PyErr_Occurred())
1125 PyErr_Clear();
1126 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001127 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001128 return NULL;
1129 }
1130 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001131 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001132 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001133 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001134 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001135 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001136 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001137 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001138 if (stop == -1) {
1139 if (PyErr_Occurred())
1140 PyErr_Clear();
1141 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001142 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001143 return NULL;
1144 }
1145 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001146 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001147 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001148 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001149 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001150 return NULL;
1151 }
1152
Raymond Hettingerb2594052004-12-05 09:25:51 +00001153 if (a3 != NULL) {
1154 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001155 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001156 if (step == -1 && PyErr_Occurred())
1157 PyErr_Clear();
1158 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001159 if (step<1) {
1160 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001161 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001162 return NULL;
1163 }
1164
1165 /* Get iterator. */
1166 it = PyObject_GetIter(seq);
1167 if (it == NULL)
1168 return NULL;
1169
1170 /* create isliceobject structure */
1171 lz = (isliceobject *)type->tp_alloc(type, 0);
1172 if (lz == NULL) {
1173 Py_DECREF(it);
1174 return NULL;
1175 }
1176 lz->it = it;
1177 lz->next = start;
1178 lz->stop = stop;
1179 lz->step = step;
1180 lz->cnt = 0L;
1181
1182 return (PyObject *)lz;
1183}
1184
1185static void
1186islice_dealloc(isliceobject *lz)
1187{
1188 PyObject_GC_UnTrack(lz);
1189 Py_XDECREF(lz->it);
1190 lz->ob_type->tp_free(lz);
1191}
1192
1193static int
1194islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1195{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001196 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001197 return 0;
1198}
1199
1200static PyObject *
1201islice_next(isliceobject *lz)
1202{
1203 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001204 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001205 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001206 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001208 assert(PyIter_Check(it));
1209 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001211 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001212 if (item == NULL)
1213 return NULL;
1214 Py_DECREF(item);
1215 lz->cnt++;
1216 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001217 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001218 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001219 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001220 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221 if (item == NULL)
1222 return NULL;
1223 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001224 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001225 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001226 if (lz->next < oldnext) /* Check for overflow */
1227 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228 return item;
1229}
1230
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231PyDoc_STRVAR(islice_doc,
1232"islice(iterable, [start,] stop [, step]) --> islice object\n\
1233\n\
1234Return an iterator whose next() method returns selected values from an\n\
1235iterable. If start is specified, will skip all preceding elements;\n\
1236otherwise, start defaults to zero. Step defaults to one. If\n\
1237specified as another value, step determines how many values are \n\
1238skipped between successive calls. Works like a slice() on a list\n\
1239but returns an iterator.");
1240
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001241static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001242 PyObject_HEAD_INIT(NULL)
1243 0, /* ob_size */
1244 "itertools.islice", /* tp_name */
1245 sizeof(isliceobject), /* tp_basicsize */
1246 0, /* tp_itemsize */
1247 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001248 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001249 0, /* tp_print */
1250 0, /* tp_getattr */
1251 0, /* tp_setattr */
1252 0, /* tp_compare */
1253 0, /* tp_repr */
1254 0, /* tp_as_number */
1255 0, /* tp_as_sequence */
1256 0, /* tp_as_mapping */
1257 0, /* tp_hash */
1258 0, /* tp_call */
1259 0, /* tp_str */
1260 PyObject_GenericGetAttr, /* tp_getattro */
1261 0, /* tp_setattro */
1262 0, /* tp_as_buffer */
1263 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1264 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001265 islice_doc, /* tp_doc */
1266 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001267 0, /* tp_clear */
1268 0, /* tp_richcompare */
1269 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001270 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001271 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001272 0, /* tp_methods */
1273 0, /* tp_members */
1274 0, /* tp_getset */
1275 0, /* tp_base */
1276 0, /* tp_dict */
1277 0, /* tp_descr_get */
1278 0, /* tp_descr_set */
1279 0, /* tp_dictoffset */
1280 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001281 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001282 islice_new, /* tp_new */
1283 PyObject_GC_Del, /* tp_free */
1284};
1285
1286
1287/* starmap object ************************************************************/
1288
1289typedef struct {
1290 PyObject_HEAD
1291 PyObject *func;
1292 PyObject *it;
1293} starmapobject;
1294
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001295static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001296
1297static PyObject *
1298starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1299{
1300 PyObject *func, *seq;
1301 PyObject *it;
1302 starmapobject *lz;
1303
Georg Brandl02c42872005-08-26 06:42:30 +00001304 if (!_PyArg_NoKeywords("starmap()", kwds))
1305 return NULL;
1306
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1308 return NULL;
1309
1310 /* Get iterator. */
1311 it = PyObject_GetIter(seq);
1312 if (it == NULL)
1313 return NULL;
1314
1315 /* create starmapobject structure */
1316 lz = (starmapobject *)type->tp_alloc(type, 0);
1317 if (lz == NULL) {
1318 Py_DECREF(it);
1319 return NULL;
1320 }
1321 Py_INCREF(func);
1322 lz->func = func;
1323 lz->it = it;
1324
1325 return (PyObject *)lz;
1326}
1327
1328static void
1329starmap_dealloc(starmapobject *lz)
1330{
1331 PyObject_GC_UnTrack(lz);
1332 Py_XDECREF(lz->func);
1333 Py_XDECREF(lz->it);
1334 lz->ob_type->tp_free(lz);
1335}
1336
1337static int
1338starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1339{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001340 Py_VISIT(lz->it);
1341 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342 return 0;
1343}
1344
1345static PyObject *
1346starmap_next(starmapobject *lz)
1347{
1348 PyObject *args;
1349 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001350 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001352 assert(PyIter_Check(it));
1353 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001354 if (args == NULL)
1355 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001356 if (!PyTuple_CheckExact(args)) {
1357 Py_DECREF(args);
1358 PyErr_SetString(PyExc_TypeError,
1359 "iterator must return a tuple");
1360 return NULL;
1361 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362 result = PyObject_Call(lz->func, args, NULL);
1363 Py_DECREF(args);
1364 return result;
1365}
1366
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367PyDoc_STRVAR(starmap_doc,
1368"starmap(function, sequence) --> starmap object\n\
1369\n\
1370Return an iterator whose values are returned from the function evaluated\n\
1371with a argument tuple taken from the given sequence.");
1372
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001373static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374 PyObject_HEAD_INIT(NULL)
1375 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001376 "itertools.starmap", /* tp_name */
1377 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378 0, /* tp_itemsize */
1379 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001380 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001381 0, /* tp_print */
1382 0, /* tp_getattr */
1383 0, /* tp_setattr */
1384 0, /* tp_compare */
1385 0, /* tp_repr */
1386 0, /* tp_as_number */
1387 0, /* tp_as_sequence */
1388 0, /* tp_as_mapping */
1389 0, /* tp_hash */
1390 0, /* tp_call */
1391 0, /* tp_str */
1392 PyObject_GenericGetAttr, /* tp_getattro */
1393 0, /* tp_setattro */
1394 0, /* tp_as_buffer */
1395 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1396 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001397 starmap_doc, /* tp_doc */
1398 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001399 0, /* tp_clear */
1400 0, /* tp_richcompare */
1401 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001402 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001403 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001404 0, /* tp_methods */
1405 0, /* tp_members */
1406 0, /* tp_getset */
1407 0, /* tp_base */
1408 0, /* tp_dict */
1409 0, /* tp_descr_get */
1410 0, /* tp_descr_set */
1411 0, /* tp_dictoffset */
1412 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001413 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001414 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001415 PyObject_GC_Del, /* tp_free */
1416};
1417
1418
1419/* imap object ************************************************************/
1420
1421typedef struct {
1422 PyObject_HEAD
1423 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001424 PyObject *func;
1425} imapobject;
1426
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001427static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001428
1429static PyObject *
1430imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1431{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001432 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001433 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001434 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435
Georg Brandl02c42872005-08-26 06:42:30 +00001436 if (!_PyArg_NoKeywords("imap()", kwds))
1437 return NULL;
1438
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001439 numargs = PyTuple_Size(args);
1440 if (numargs < 2) {
1441 PyErr_SetString(PyExc_TypeError,
1442 "imap() must have at least two arguments.");
1443 return NULL;
1444 }
1445
1446 iters = PyTuple_New(numargs-1);
1447 if (iters == NULL)
1448 return NULL;
1449
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001450 for (i=1 ; i<numargs ; i++) {
1451 /* Get iterator. */
1452 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1453 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454 Py_DECREF(iters);
1455 return NULL;
1456 }
1457 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001458 }
1459
1460 /* create imapobject structure */
1461 lz = (imapobject *)type->tp_alloc(type, 0);
1462 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001463 Py_DECREF(iters);
1464 return NULL;
1465 }
1466 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467 func = PyTuple_GET_ITEM(args, 0);
1468 Py_INCREF(func);
1469 lz->func = func;
1470
1471 return (PyObject *)lz;
1472}
1473
1474static void
1475imap_dealloc(imapobject *lz)
1476{
1477 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001478 Py_XDECREF(lz->iters);
1479 Py_XDECREF(lz->func);
1480 lz->ob_type->tp_free(lz);
1481}
1482
1483static int
1484imap_traverse(imapobject *lz, visitproc visit, void *arg)
1485{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001486 Py_VISIT(lz->iters);
1487 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001488 return 0;
1489}
1490
Raymond Hettinger2012f172003-02-07 05:32:58 +00001491/*
1492imap() is an iterator version of __builtins__.map() except that it does
1493not have the None fill-in feature. That was intentionally left out for
1494the following reasons:
1495
1496 1) Itertools are designed to be easily combined and chained together.
1497 Having all tools stop with the shortest input is a unifying principle
1498 that makes it easier to combine finite iterators (supplying data) with
1499 infinite iterators like count() and repeat() (for supplying sequential
1500 or constant arguments to a function).
1501
1502 2) In typical use cases for combining itertools, having one finite data
1503 supplier run out before another is likely to be an error condition which
1504 should not pass silently by automatically supplying None.
1505
1506 3) The use cases for automatic None fill-in are rare -- not many functions
1507 do something useful when a parameter suddenly switches type and becomes
1508 None.
1509
1510 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001511 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001512
1513 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1514*/
1515
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001516static PyObject *
1517imap_next(imapobject *lz)
1518{
1519 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001520 PyObject *argtuple;
1521 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001522 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523
1524 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001525 argtuple = PyTuple_New(numargs);
1526 if (argtuple == NULL)
1527 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001528
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001529 for (i=0 ; i<numargs ; i++) {
1530 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1531 if (val == NULL) {
1532 Py_DECREF(argtuple);
1533 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001534 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001535 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001537 if (lz->func == Py_None)
1538 return argtuple;
1539 result = PyObject_Call(lz->func, argtuple, NULL);
1540 Py_DECREF(argtuple);
1541 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001542}
1543
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001544PyDoc_STRVAR(imap_doc,
1545"imap(func, *iterables) --> imap object\n\
1546\n\
1547Make an iterator that computes the function using arguments from\n\
1548each of the iterables. Like map() except that it returns\n\
1549an iterator instead of a list and that it stops when the shortest\n\
1550iterable is exhausted instead of filling in None for shorter\n\
1551iterables.");
1552
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001553static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554 PyObject_HEAD_INIT(NULL)
1555 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001556 "itertools.imap", /* tp_name */
1557 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001558 0, /* tp_itemsize */
1559 /* methods */
1560 (destructor)imap_dealloc, /* tp_dealloc */
1561 0, /* tp_print */
1562 0, /* tp_getattr */
1563 0, /* tp_setattr */
1564 0, /* tp_compare */
1565 0, /* tp_repr */
1566 0, /* tp_as_number */
1567 0, /* tp_as_sequence */
1568 0, /* tp_as_mapping */
1569 0, /* tp_hash */
1570 0, /* tp_call */
1571 0, /* tp_str */
1572 PyObject_GenericGetAttr, /* tp_getattro */
1573 0, /* tp_setattro */
1574 0, /* tp_as_buffer */
1575 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1576 Py_TPFLAGS_BASETYPE, /* tp_flags */
1577 imap_doc, /* tp_doc */
1578 (traverseproc)imap_traverse, /* tp_traverse */
1579 0, /* tp_clear */
1580 0, /* tp_richcompare */
1581 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001582 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583 (iternextfunc)imap_next, /* tp_iternext */
1584 0, /* tp_methods */
1585 0, /* tp_members */
1586 0, /* tp_getset */
1587 0, /* tp_base */
1588 0, /* tp_dict */
1589 0, /* tp_descr_get */
1590 0, /* tp_descr_set */
1591 0, /* tp_dictoffset */
1592 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001593 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001594 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595 PyObject_GC_Del, /* tp_free */
1596};
1597
1598
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001599/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001600
1601typedef struct {
1602 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00001603 Py_ssize_t tuplesize;
1604 Py_ssize_t iternum; /* which iterator is active */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001605 PyObject *ittuple; /* tuple of iterators */
1606} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001608static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609
1610static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001611chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001612{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001613 chainobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001614 Py_ssize_t tuplesize = PySequence_Length(args);
Jack Diederich6c433a92006-05-26 11:15:17 +00001615 Py_ssize_t i;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001616 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617
Georg Brandl02c42872005-08-26 06:42:30 +00001618 if (!_PyArg_NoKeywords("chain()", kwds))
1619 return NULL;
1620
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001621 /* obtain iterators */
1622 assert(PyTuple_Check(args));
1623 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001624 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001625 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001626 for (i=0; i < tuplesize; ++i) {
1627 PyObject *item = PyTuple_GET_ITEM(args, i);
1628 PyObject *it = PyObject_GetIter(item);
1629 if (it == NULL) {
1630 if (PyErr_ExceptionMatches(PyExc_TypeError))
1631 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001632 "chain argument #%zd must support iteration",
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001633 i+1);
1634 Py_DECREF(ittuple);
1635 return NULL;
1636 }
1637 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638 }
1639
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001640 /* create chainobject structure */
1641 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001642 if (lz == NULL) {
1643 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001645 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001646
1647 lz->ittuple = ittuple;
1648 lz->iternum = 0;
1649 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650
1651 return (PyObject *)lz;
1652}
1653
1654static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001655chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001656{
1657 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001658 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001659 lz->ob_type->tp_free(lz);
1660}
1661
Raymond Hettinger2012f172003-02-07 05:32:58 +00001662static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001663chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001664{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001665 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001666 return 0;
1667}
1668
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001669static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001670chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001672 PyObject *it;
1673 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001675 while (lz->iternum < lz->tuplesize) {
1676 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1677 item = PyIter_Next(it);
1678 if (item != NULL)
1679 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001680 if (PyErr_Occurred()) {
1681 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1682 PyErr_Clear();
1683 else
1684 return NULL;
1685 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001686 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687 }
1688 return NULL;
1689}
1690
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001691PyDoc_STRVAR(chain_doc,
1692"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001693\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001694Return a chain object whose .next() method returns elements from the\n\
1695first iterable until it is exhausted, then elements from the next\n\
1696iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001698static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001699 PyObject_HEAD_INIT(NULL)
1700 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001701 "itertools.chain", /* tp_name */
1702 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001703 0, /* tp_itemsize */
1704 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001705 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001706 0, /* tp_print */
1707 0, /* tp_getattr */
1708 0, /* tp_setattr */
1709 0, /* tp_compare */
1710 0, /* tp_repr */
1711 0, /* tp_as_number */
1712 0, /* tp_as_sequence */
1713 0, /* tp_as_mapping */
1714 0, /* tp_hash */
1715 0, /* tp_call */
1716 0, /* tp_str */
1717 PyObject_GenericGetAttr, /* tp_getattro */
1718 0, /* tp_setattro */
1719 0, /* tp_as_buffer */
1720 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1721 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001722 chain_doc, /* tp_doc */
1723 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724 0, /* tp_clear */
1725 0, /* tp_richcompare */
1726 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001727 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001728 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001729 0, /* tp_methods */
1730 0, /* tp_members */
1731 0, /* tp_getset */
1732 0, /* tp_base */
1733 0, /* tp_dict */
1734 0, /* tp_descr_get */
1735 0, /* tp_descr_set */
1736 0, /* tp_dictoffset */
1737 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001738 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001739 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001740 PyObject_GC_Del, /* tp_free */
1741};
1742
1743
1744/* ifilter object ************************************************************/
1745
1746typedef struct {
1747 PyObject_HEAD
1748 PyObject *func;
1749 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001750} ifilterobject;
1751
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001752static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753
1754static PyObject *
1755ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1756{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001757 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758 PyObject *it;
1759 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001760
Georg Brandl02c42872005-08-26 06:42:30 +00001761 if (!_PyArg_NoKeywords("ifilter()", kwds))
1762 return NULL;
1763
Raymond Hettinger60eca932003-02-09 06:40:58 +00001764 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001765 return NULL;
1766
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001767 /* Get iterator. */
1768 it = PyObject_GetIter(seq);
1769 if (it == NULL)
1770 return NULL;
1771
1772 /* create ifilterobject structure */
1773 lz = (ifilterobject *)type->tp_alloc(type, 0);
1774 if (lz == NULL) {
1775 Py_DECREF(it);
1776 return NULL;
1777 }
1778 Py_INCREF(func);
1779 lz->func = func;
1780 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001781
1782 return (PyObject *)lz;
1783}
1784
1785static void
1786ifilter_dealloc(ifilterobject *lz)
1787{
1788 PyObject_GC_UnTrack(lz);
1789 Py_XDECREF(lz->func);
1790 Py_XDECREF(lz->it);
1791 lz->ob_type->tp_free(lz);
1792}
1793
1794static int
1795ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1796{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001797 Py_VISIT(lz->it);
1798 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001799 return 0;
1800}
1801
1802static PyObject *
1803ifilter_next(ifilterobject *lz)
1804{
1805 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001806 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001807 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001808 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001809
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001810 assert(PyIter_Check(it));
1811 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001812 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001813 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001814 if (item == NULL)
1815 return NULL;
1816
1817 if (lz->func == Py_None) {
1818 ok = PyObject_IsTrue(item);
1819 } else {
1820 PyObject *good;
1821 good = PyObject_CallFunctionObjArgs(lz->func,
1822 item, NULL);
1823 if (good == NULL) {
1824 Py_DECREF(item);
1825 return NULL;
1826 }
1827 ok = PyObject_IsTrue(good);
1828 Py_DECREF(good);
1829 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001830 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001831 return item;
1832 Py_DECREF(item);
1833 }
1834}
1835
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001836PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001837"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001838\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001839Return those items of sequence for which function(item) is true.\n\
1840If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001841
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001842static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001843 PyObject_HEAD_INIT(NULL)
1844 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001845 "itertools.ifilter", /* tp_name */
1846 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001847 0, /* tp_itemsize */
1848 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001849 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001850 0, /* tp_print */
1851 0, /* tp_getattr */
1852 0, /* tp_setattr */
1853 0, /* tp_compare */
1854 0, /* tp_repr */
1855 0, /* tp_as_number */
1856 0, /* tp_as_sequence */
1857 0, /* tp_as_mapping */
1858 0, /* tp_hash */
1859 0, /* tp_call */
1860 0, /* tp_str */
1861 PyObject_GenericGetAttr, /* tp_getattro */
1862 0, /* tp_setattro */
1863 0, /* tp_as_buffer */
1864 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1865 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001866 ifilter_doc, /* tp_doc */
1867 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001868 0, /* tp_clear */
1869 0, /* tp_richcompare */
1870 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001871 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001872 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001873 0, /* tp_methods */
1874 0, /* tp_members */
1875 0, /* tp_getset */
1876 0, /* tp_base */
1877 0, /* tp_dict */
1878 0, /* tp_descr_get */
1879 0, /* tp_descr_set */
1880 0, /* tp_dictoffset */
1881 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001882 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001883 ifilter_new, /* tp_new */
1884 PyObject_GC_Del, /* tp_free */
1885};
1886
1887
1888/* ifilterfalse object ************************************************************/
1889
1890typedef struct {
1891 PyObject_HEAD
1892 PyObject *func;
1893 PyObject *it;
1894} ifilterfalseobject;
1895
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001896static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001897
1898static PyObject *
1899ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1900{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001901 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001902 PyObject *it;
1903 ifilterfalseobject *lz;
1904
Georg Brandl02c42872005-08-26 06:42:30 +00001905 if (!_PyArg_NoKeywords("ifilterfalse()", kwds))
1906 return NULL;
1907
Raymond Hettinger60eca932003-02-09 06:40:58 +00001908 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1909 return NULL;
1910
1911 /* Get iterator. */
1912 it = PyObject_GetIter(seq);
1913 if (it == NULL)
1914 return NULL;
1915
1916 /* create ifilterfalseobject structure */
1917 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1918 if (lz == NULL) {
1919 Py_DECREF(it);
1920 return NULL;
1921 }
1922 Py_INCREF(func);
1923 lz->func = func;
1924 lz->it = it;
1925
1926 return (PyObject *)lz;
1927}
1928
1929static void
1930ifilterfalse_dealloc(ifilterfalseobject *lz)
1931{
1932 PyObject_GC_UnTrack(lz);
1933 Py_XDECREF(lz->func);
1934 Py_XDECREF(lz->it);
1935 lz->ob_type->tp_free(lz);
1936}
1937
1938static int
1939ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1940{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001941 Py_VISIT(lz->it);
1942 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001943 return 0;
1944}
1945
1946static PyObject *
1947ifilterfalse_next(ifilterfalseobject *lz)
1948{
1949 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001950 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001951 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001952 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001953
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001954 assert(PyIter_Check(it));
1955 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001956 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001957 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001958 if (item == NULL)
1959 return NULL;
1960
1961 if (lz->func == Py_None) {
1962 ok = PyObject_IsTrue(item);
1963 } else {
1964 PyObject *good;
1965 good = PyObject_CallFunctionObjArgs(lz->func,
1966 item, NULL);
1967 if (good == NULL) {
1968 Py_DECREF(item);
1969 return NULL;
1970 }
1971 ok = PyObject_IsTrue(good);
1972 Py_DECREF(good);
1973 }
1974 if (!ok)
1975 return item;
1976 Py_DECREF(item);
1977 }
1978}
1979
Raymond Hettinger60eca932003-02-09 06:40:58 +00001980PyDoc_STRVAR(ifilterfalse_doc,
1981"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1982\n\
1983Return those items of sequence for which function(item) is false.\n\
1984If function is None, return the items that are false.");
1985
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001986static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001987 PyObject_HEAD_INIT(NULL)
1988 0, /* ob_size */
1989 "itertools.ifilterfalse", /* tp_name */
1990 sizeof(ifilterfalseobject), /* tp_basicsize */
1991 0, /* tp_itemsize */
1992 /* methods */
1993 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1994 0, /* tp_print */
1995 0, /* tp_getattr */
1996 0, /* tp_setattr */
1997 0, /* tp_compare */
1998 0, /* tp_repr */
1999 0, /* tp_as_number */
2000 0, /* tp_as_sequence */
2001 0, /* tp_as_mapping */
2002 0, /* tp_hash */
2003 0, /* tp_call */
2004 0, /* tp_str */
2005 PyObject_GenericGetAttr, /* tp_getattro */
2006 0, /* tp_setattro */
2007 0, /* tp_as_buffer */
2008 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2009 Py_TPFLAGS_BASETYPE, /* tp_flags */
2010 ifilterfalse_doc, /* tp_doc */
2011 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2012 0, /* tp_clear */
2013 0, /* tp_richcompare */
2014 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002015 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002016 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2017 0, /* tp_methods */
2018 0, /* tp_members */
2019 0, /* tp_getset */
2020 0, /* tp_base */
2021 0, /* tp_dict */
2022 0, /* tp_descr_get */
2023 0, /* tp_descr_set */
2024 0, /* tp_dictoffset */
2025 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002026 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002027 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002028 PyObject_GC_Del, /* tp_free */
2029};
2030
2031
2032/* count object ************************************************************/
2033
2034typedef struct {
2035 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002036 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002037} countobject;
2038
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002039static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002040
2041static PyObject *
2042count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2043{
2044 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002045 Py_ssize_t cnt = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002046
Georg Brandl02c42872005-08-26 06:42:30 +00002047 if (!_PyArg_NoKeywords("count()", kwds))
2048 return NULL;
2049
Jack Diederich6c433a92006-05-26 11:15:17 +00002050 if (!PyArg_ParseTuple(args, "|n:count", &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002051 return NULL;
2052
2053 /* create countobject structure */
2054 lz = (countobject *)PyObject_New(countobject, &count_type);
2055 if (lz == NULL)
2056 return NULL;
2057 lz->cnt = cnt;
2058
2059 return (PyObject *)lz;
2060}
2061
2062static PyObject *
2063count_next(countobject *lz)
2064{
Jack Diederich6c433a92006-05-26 11:15:17 +00002065 return PyInt_FromSize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002066}
2067
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002068static PyObject *
2069count_repr(countobject *lz)
2070{
Jack Diederich6c433a92006-05-26 11:15:17 +00002071 return PyString_FromFormat("count(%zd)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002072}
2073
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002074PyDoc_STRVAR(count_doc,
2075"count([firstval]) --> count object\n\
2076\n\
2077Return a count object whose .next() method returns consecutive\n\
2078integers starting from zero or, if specified, from firstval.");
2079
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002080static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002081 PyObject_HEAD_INIT(NULL)
2082 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002083 "itertools.count", /* tp_name */
2084 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002085 0, /* tp_itemsize */
2086 /* methods */
2087 (destructor)PyObject_Del, /* tp_dealloc */
2088 0, /* tp_print */
2089 0, /* tp_getattr */
2090 0, /* tp_setattr */
2091 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002092 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002093 0, /* tp_as_number */
2094 0, /* tp_as_sequence */
2095 0, /* tp_as_mapping */
2096 0, /* tp_hash */
2097 0, /* tp_call */
2098 0, /* tp_str */
2099 PyObject_GenericGetAttr, /* tp_getattro */
2100 0, /* tp_setattro */
2101 0, /* tp_as_buffer */
2102 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002103 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002104 0, /* tp_traverse */
2105 0, /* tp_clear */
2106 0, /* tp_richcompare */
2107 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002108 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002109 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002110 0, /* tp_methods */
2111 0, /* tp_members */
2112 0, /* tp_getset */
2113 0, /* tp_base */
2114 0, /* tp_dict */
2115 0, /* tp_descr_get */
2116 0, /* tp_descr_set */
2117 0, /* tp_dictoffset */
2118 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002119 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002120 count_new, /* tp_new */
2121};
2122
2123
2124/* izip object ************************************************************/
2125
2126#include "Python.h"
2127
2128typedef struct {
2129 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002130 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002131 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002132 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002133} izipobject;
2134
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002135static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002136
2137static PyObject *
2138izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2139{
2140 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002141 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002142 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002143 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002144 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002145
Georg Brandl02c42872005-08-26 06:42:30 +00002146 if (!_PyArg_NoKeywords("izip()", kwds))
2147 return NULL;
2148
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002149 /* args must be a tuple */
2150 assert(PyTuple_Check(args));
2151
2152 /* obtain iterators */
2153 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002154 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002155 return NULL;
2156 for (i=0; i < tuplesize; ++i) {
2157 PyObject *item = PyTuple_GET_ITEM(args, i);
2158 PyObject *it = PyObject_GetIter(item);
2159 if (it == NULL) {
2160 if (PyErr_ExceptionMatches(PyExc_TypeError))
2161 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002162 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002163 i+1);
2164 Py_DECREF(ittuple);
2165 return NULL;
2166 }
2167 PyTuple_SET_ITEM(ittuple, i, it);
2168 }
2169
Raymond Hettinger2012f172003-02-07 05:32:58 +00002170 /* create a result holder */
2171 result = PyTuple_New(tuplesize);
2172 if (result == NULL) {
2173 Py_DECREF(ittuple);
2174 return NULL;
2175 }
2176 for (i=0 ; i < tuplesize ; i++) {
2177 Py_INCREF(Py_None);
2178 PyTuple_SET_ITEM(result, i, Py_None);
2179 }
2180
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002181 /* create izipobject structure */
2182 lz = (izipobject *)type->tp_alloc(type, 0);
2183 if (lz == NULL) {
2184 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002185 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002186 return NULL;
2187 }
2188 lz->ittuple = ittuple;
2189 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002190 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002191
2192 return (PyObject *)lz;
2193}
2194
2195static void
2196izip_dealloc(izipobject *lz)
2197{
2198 PyObject_GC_UnTrack(lz);
2199 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002200 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002201 lz->ob_type->tp_free(lz);
2202}
2203
2204static int
2205izip_traverse(izipobject *lz, visitproc visit, void *arg)
2206{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002207 Py_VISIT(lz->ittuple);
2208 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002209 return 0;
2210}
2211
2212static PyObject *
2213izip_next(izipobject *lz)
2214{
Jack Diederich6c433a92006-05-26 11:15:17 +00002215 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002216 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002217 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002218 PyObject *it;
2219 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002220 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002221
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002222 if (tuplesize == 0)
2223 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002224 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002225 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002226 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002227 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002228 assert(PyIter_Check(it));
2229 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002230 if (item == NULL) {
2231 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002232 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002233 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002234 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002235 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002236 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002237 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002238 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002239 result = PyTuple_New(tuplesize);
2240 if (result == NULL)
2241 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002242 for (i=0 ; i < tuplesize ; i++) {
2243 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002244 assert(PyIter_Check(it));
2245 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002246 if (item == NULL) {
2247 Py_DECREF(result);
2248 return NULL;
2249 }
2250 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002251 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002252 }
2253 return result;
2254}
2255
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002256PyDoc_STRVAR(izip_doc,
2257"izip(iter1 [,iter2 [...]]) --> izip object\n\
2258\n\
2259Return a izip object whose .next() method returns a tuple where\n\
2260the i-th element comes from the i-th iterable argument. The .next()\n\
2261method continues until the shortest iterable in the argument sequence\n\
2262is exhausted and then it raises StopIteration. Works like the zip()\n\
2263function but consumes less memory by returning an iterator instead of\n\
2264a list.");
2265
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002266static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002267 PyObject_HEAD_INIT(NULL)
2268 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002269 "itertools.izip", /* tp_name */
2270 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002271 0, /* tp_itemsize */
2272 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002273 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002274 0, /* tp_print */
2275 0, /* tp_getattr */
2276 0, /* tp_setattr */
2277 0, /* tp_compare */
2278 0, /* tp_repr */
2279 0, /* tp_as_number */
2280 0, /* tp_as_sequence */
2281 0, /* tp_as_mapping */
2282 0, /* tp_hash */
2283 0, /* tp_call */
2284 0, /* tp_str */
2285 PyObject_GenericGetAttr, /* tp_getattro */
2286 0, /* tp_setattro */
2287 0, /* tp_as_buffer */
2288 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2289 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002290 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002291 (traverseproc)izip_traverse, /* tp_traverse */
2292 0, /* tp_clear */
2293 0, /* tp_richcompare */
2294 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002295 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002296 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002297 0, /* tp_methods */
2298 0, /* tp_members */
2299 0, /* tp_getset */
2300 0, /* tp_base */
2301 0, /* tp_dict */
2302 0, /* tp_descr_get */
2303 0, /* tp_descr_set */
2304 0, /* tp_dictoffset */
2305 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002306 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002307 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002308 PyObject_GC_Del, /* tp_free */
2309};
2310
2311
2312/* repeat object ************************************************************/
2313
2314typedef struct {
2315 PyObject_HEAD
2316 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002317 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002318} repeatobject;
2319
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002320static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002321
2322static PyObject *
2323repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2324{
2325 repeatobject *ro;
2326 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002327 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002328
Georg Brandl02c42872005-08-26 06:42:30 +00002329 if (!_PyArg_NoKeywords("repeat()", kwds))
2330 return NULL;
2331
Jack Diederich6c433a92006-05-26 11:15:17 +00002332 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002333 return NULL;
2334
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002335 if (PyTuple_Size(args) == 2 && cnt < 0)
2336 cnt = 0;
2337
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002338 ro = (repeatobject *)type->tp_alloc(type, 0);
2339 if (ro == NULL)
2340 return NULL;
2341 Py_INCREF(element);
2342 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002343 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002344 return (PyObject *)ro;
2345}
2346
2347static void
2348repeat_dealloc(repeatobject *ro)
2349{
2350 PyObject_GC_UnTrack(ro);
2351 Py_XDECREF(ro->element);
2352 ro->ob_type->tp_free(ro);
2353}
2354
2355static int
2356repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2357{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002358 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002359 return 0;
2360}
2361
2362static PyObject *
2363repeat_next(repeatobject *ro)
2364{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002365 if (ro->cnt == 0)
2366 return NULL;
2367 if (ro->cnt > 0)
2368 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002369 Py_INCREF(ro->element);
2370 return ro->element;
2371}
2372
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002373static PyObject *
2374repeat_repr(repeatobject *ro)
2375{
2376 PyObject *result, *objrepr;
2377
2378 objrepr = PyObject_Repr(ro->element);
2379 if (objrepr == NULL)
2380 return NULL;
2381
2382 if (ro->cnt == -1)
2383 result = PyString_FromFormat("repeat(%s)",
2384 PyString_AS_STRING(objrepr));
2385 else
Jack Diederich6c433a92006-05-26 11:15:17 +00002386 result = PyString_FromFormat("repeat(%s, %zd)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002387 PyString_AS_STRING(objrepr), ro->cnt);
2388 Py_DECREF(objrepr);
2389 return result;
2390}
2391
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002392static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002393repeat_len(repeatobject *ro)
2394{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002395 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002396 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002397 return NULL;
2398 }
Jack Diederich6c433a92006-05-26 11:15:17 +00002399 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002400}
2401
Armin Rigof5b3e362006-02-11 21:32:43 +00002402PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002403
2404static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002405 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002406 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002407};
2408
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002409PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002410"repeat(element [,times]) -> create an iterator which returns the element\n\
2411for the specified number of times. If not specified, returns the element\n\
2412endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002413
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002414static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002415 PyObject_HEAD_INIT(NULL)
2416 0, /* ob_size */
2417 "itertools.repeat", /* tp_name */
2418 sizeof(repeatobject), /* tp_basicsize */
2419 0, /* tp_itemsize */
2420 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002421 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002422 0, /* tp_print */
2423 0, /* tp_getattr */
2424 0, /* tp_setattr */
2425 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002426 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002427 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002428 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002429 0, /* tp_as_mapping */
2430 0, /* tp_hash */
2431 0, /* tp_call */
2432 0, /* tp_str */
2433 PyObject_GenericGetAttr, /* tp_getattro */
2434 0, /* tp_setattro */
2435 0, /* tp_as_buffer */
2436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2437 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002438 repeat_doc, /* tp_doc */
2439 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002440 0, /* tp_clear */
2441 0, /* tp_richcompare */
2442 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002443 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002444 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002445 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002446 0, /* tp_members */
2447 0, /* tp_getset */
2448 0, /* tp_base */
2449 0, /* tp_dict */
2450 0, /* tp_descr_get */
2451 0, /* tp_descr_set */
2452 0, /* tp_dictoffset */
2453 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002454 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002455 repeat_new, /* tp_new */
2456 PyObject_GC_Del, /* tp_free */
2457};
2458
2459
2460/* module level code ********************************************************/
2461
2462PyDoc_STRVAR(module_doc,
2463"Functional tools for creating and using iterators.\n\
2464\n\
2465Infinite iterators:\n\
2466count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002467cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002468repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002469\n\
2470Iterators terminating on the shortest input sequence:\n\
2471izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002472ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2473ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002474islice(seq, [start,] stop [, step]) --> elements from\n\
2475 seq[start:stop:step]\n\
2476imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2477starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002478tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002479chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002480takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2481dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002482groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002483");
2484
2485
Raymond Hettingerad983e72003-11-12 14:32:26 +00002486static PyMethodDef module_methods[] = {
2487 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2488 {NULL, NULL} /* sentinel */
2489};
2490
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002491PyMODINIT_FUNC
2492inititertools(void)
2493{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002494 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002495 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002496 char *name;
2497 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002498 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002499 &dropwhile_type,
2500 &takewhile_type,
2501 &islice_type,
2502 &starmap_type,
2503 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002504 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002505 &ifilter_type,
2506 &ifilterfalse_type,
2507 &count_type,
2508 &izip_type,
2509 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002510 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002511 NULL
2512 };
2513
Skip Montanarof3938fd2004-02-10 20:27:40 +00002514 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002515 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002516 if (m == NULL)
2517 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002518
Raymond Hettinger60eca932003-02-09 06:40:58 +00002519 for (i=0 ; typelist[i] != NULL ; i++) {
2520 if (PyType_Ready(typelist[i]) < 0)
2521 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002522 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002523 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002524 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002525 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002526 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002527
2528 if (PyType_Ready(&teedataobject_type) < 0)
2529 return;
2530 if (PyType_Ready(&tee_type) < 0)
2531 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002532 if (PyType_Ready(&_grouper_type) < 0)
2533 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002534}