blob: 387133c474ab1e0b18331ed26c2aaf7211848db7 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
3
4/* Itertools module written and maintained
5 by Raymond D. Hettinger <python@rcn.com>
6 Copyright (c) 2003 Python Software Foundation.
7 All rights reserved.
8*/
9
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000010
11/* groupby object ***********************************************************/
12
13typedef struct {
14 PyObject_HEAD
15 PyObject *it;
16 PyObject *keyfunc;
17 PyObject *tgtkey;
18 PyObject *currkey;
19 PyObject *currvalue;
20} groupbyobject;
21
22static PyTypeObject groupby_type;
23static PyObject *_grouper_create(groupbyobject *, PyObject *);
24
25static PyObject *
26groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
27{
28 static char *kwargs[] = {"iterable", "key", NULL};
29 groupbyobject *gbo;
30 PyObject *it, *keyfunc = Py_None;
31
32 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
33 &it, &keyfunc))
34 return NULL;
35
36 gbo = (groupbyobject *)type->tp_alloc(type, 0);
37 if (gbo == NULL)
38 return NULL;
39 gbo->tgtkey = NULL;
40 gbo->currkey = NULL;
41 gbo->currvalue = NULL;
42 gbo->keyfunc = keyfunc;
43 Py_INCREF(keyfunc);
44 gbo->it = PyObject_GetIter(it);
45 if (gbo->it == NULL) {
46 Py_DECREF(gbo);
47 return NULL;
48 }
49 return (PyObject *)gbo;
50}
51
52static void
53groupby_dealloc(groupbyobject *gbo)
54{
55 PyObject_GC_UnTrack(gbo);
56 Py_XDECREF(gbo->it);
57 Py_XDECREF(gbo->keyfunc);
58 Py_XDECREF(gbo->tgtkey);
59 Py_XDECREF(gbo->currkey);
60 Py_XDECREF(gbo->currvalue);
61 gbo->ob_type->tp_free(gbo);
62}
63
64static int
65groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
66{
67 int err;
68
69 if (gbo->it) {
70 err = visit(gbo->it, arg);
71 if (err)
72 return err;
73 }
74 if (gbo->keyfunc) {
75 err = visit(gbo->keyfunc, arg);
76 if (err)
77 return err;
78 }
79 if (gbo->tgtkey) {
80 err = visit(gbo->tgtkey, arg);
81 if (err)
82 return err;
83 }
84 if (gbo->currkey) {
85 err = visit(gbo->currkey, arg);
86 if (err)
87 return err;
88 }
89 if (gbo->currvalue) {
90 err = visit(gbo->currvalue, arg);
91 if (err)
92 return err;
93 }
94 return 0;
95}
96
97static PyObject *
98groupby_next(groupbyobject *gbo)
99{
100 PyObject *newvalue, *newkey, *r, *grouper;
101
102 /* skip to next iteration group */
103 for (;;) {
104 if (gbo->currkey == NULL)
105 /* pass */;
106 else if (gbo->tgtkey == NULL)
107 break;
108 else {
109 int rcmp;
110
111 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
112 gbo->currkey, Py_EQ);
113 if (rcmp == -1)
114 return NULL;
115 else if (rcmp == 0)
116 break;
117 }
118
119 newvalue = PyIter_Next(gbo->it);
120 if (newvalue == NULL)
121 return NULL;
122
123 if (gbo->keyfunc == Py_None) {
124 newkey = newvalue;
125 Py_INCREF(newvalue);
126 } else {
127 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
128 newvalue, NULL);
129 if (newkey == NULL) {
130 Py_DECREF(newvalue);
131 return NULL;
132 }
133 }
134
135 Py_XDECREF(gbo->currkey);
136 gbo->currkey = newkey;
137 Py_XDECREF(gbo->currvalue);
138 gbo->currvalue = newvalue;
139 }
140
141 Py_XDECREF(gbo->tgtkey);
142 gbo->tgtkey = gbo->currkey;
143 Py_INCREF(gbo->currkey);
144
145 grouper = _grouper_create(gbo, gbo->tgtkey);
146 if (grouper == NULL)
147 return NULL;
148
149 r = PyTuple_Pack(2, gbo->currkey, grouper);
150 Py_DECREF(grouper);
151 return r;
152}
153
154PyDoc_STRVAR(groupby_doc,
155"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
156(key, sub-iterator) grouped by each value of key(value).\n");
157
158static PyTypeObject groupby_type = {
159 PyObject_HEAD_INIT(NULL)
160 0, /* ob_size */
161 "itertools.groupby", /* tp_name */
162 sizeof(groupbyobject), /* tp_basicsize */
163 0, /* tp_itemsize */
164 /* methods */
165 (destructor)groupby_dealloc, /* tp_dealloc */
166 0, /* tp_print */
167 0, /* tp_getattr */
168 0, /* tp_setattr */
169 0, /* tp_compare */
170 0, /* tp_repr */
171 0, /* tp_as_number */
172 0, /* tp_as_sequence */
173 0, /* tp_as_mapping */
174 0, /* tp_hash */
175 0, /* tp_call */
176 0, /* tp_str */
177 PyObject_GenericGetAttr, /* tp_getattro */
178 0, /* tp_setattro */
179 0, /* tp_as_buffer */
180 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
181 Py_TPFLAGS_BASETYPE, /* tp_flags */
182 groupby_doc, /* tp_doc */
183 (traverseproc)groupby_traverse, /* tp_traverse */
184 0, /* tp_clear */
185 0, /* tp_richcompare */
186 0, /* tp_weaklistoffset */
187 PyObject_SelfIter, /* tp_iter */
188 (iternextfunc)groupby_next, /* tp_iternext */
189 0, /* tp_methods */
190 0, /* tp_members */
191 0, /* tp_getset */
192 0, /* tp_base */
193 0, /* tp_dict */
194 0, /* tp_descr_get */
195 0, /* tp_descr_set */
196 0, /* tp_dictoffset */
197 0, /* tp_init */
198 0, /* tp_alloc */
199 groupby_new, /* tp_new */
200 PyObject_GC_Del, /* tp_free */
201};
202
203
204/* _grouper object (internal) ************************************************/
205
206typedef struct {
207 PyObject_HEAD
208 PyObject *parent;
209 PyObject *tgtkey;
210} _grouperobject;
211
212static PyTypeObject _grouper_type;
213
214static PyObject *
215_grouper_create(groupbyobject *parent, PyObject *tgtkey)
216{
217 _grouperobject *igo;
218
219 igo = PyObject_New(_grouperobject, &_grouper_type);
220 if (igo == NULL)
221 return NULL;
222 igo->parent = (PyObject *)parent;
223 Py_INCREF(parent);
224 igo->tgtkey = tgtkey;
225 Py_INCREF(tgtkey);
226
227 return (PyObject *)igo;
228}
229
230static void
231_grouper_dealloc(_grouperobject *igo)
232{
233 Py_DECREF(igo->parent);
234 Py_DECREF(igo->tgtkey);
235 PyObject_Del(igo);
236}
237
238static PyObject *
239_grouper_next(_grouperobject *igo)
240{
241 groupbyobject *gbo = (groupbyobject *)igo->parent;
242 PyObject *newvalue, *newkey, *r;
243 int rcmp;
244
245 if (gbo->currvalue == NULL) {
246 newvalue = PyIter_Next(gbo->it);
247 if (newvalue == NULL)
248 return NULL;
249
250 if (gbo->keyfunc == Py_None) {
251 newkey = newvalue;
252 Py_INCREF(newvalue);
253 } else {
254 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
255 newvalue, NULL);
256 if (newkey == NULL) {
257 Py_DECREF(newvalue);
258 return NULL;
259 }
260 }
261
262 assert(gbo->currkey == NULL);
263 gbo->currkey = newkey;
264 gbo->currvalue = newvalue;
265 }
266
267 assert(gbo->currkey != NULL);
268 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
269 if (rcmp <= 0)
270 /* got any error or current group is end */
271 return NULL;
272
273 r = gbo->currvalue;
274 gbo->currvalue = NULL;
275 Py_DECREF(gbo->currkey);
276 gbo->currkey = NULL;
277
278 return r;
279}
280
281static PyTypeObject _grouper_type = {
282 PyObject_HEAD_INIT(NULL)
283 0, /* ob_size */
284 "itertools._grouper", /* tp_name */
285 sizeof(_grouperobject), /* tp_basicsize */
286 0, /* tp_itemsize */
287 /* methods */
288 (destructor)_grouper_dealloc, /* tp_dealloc */
289 0, /* tp_print */
290 0, /* tp_getattr */
291 0, /* tp_setattr */
292 0, /* tp_compare */
293 0, /* tp_repr */
294 0, /* tp_as_number */
295 0, /* tp_as_sequence */
296 0, /* tp_as_mapping */
297 0, /* tp_hash */
298 0, /* tp_call */
299 0, /* tp_str */
300 PyObject_GenericGetAttr, /* tp_getattro */
301 0, /* tp_setattro */
302 0, /* tp_as_buffer */
303 Py_TPFLAGS_DEFAULT, /* tp_flags */
304 0, /* tp_doc */
305 0, /* tp_traverse */
306 0, /* tp_clear */
307 0, /* tp_richcompare */
308 0, /* tp_weaklistoffset */
309 PyObject_SelfIter, /* tp_iter */
310 (iternextfunc)_grouper_next, /* tp_iternext */
311 0, /* tp_methods */
312 0, /* tp_members */
313 0, /* tp_getset */
314 0, /* tp_base */
315 0, /* tp_dict */
316 0, /* tp_descr_get */
317 0, /* tp_descr_set */
318 0, /* tp_dictoffset */
319 0, /* tp_init */
320 0, /* tp_alloc */
321 0, /* tp_new */
322 PyObject_Del, /* tp_free */
323};
324
325
326
Raymond Hettingerad983e72003-11-12 14:32:26 +0000327/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000328
Raymond Hettingerad983e72003-11-12 14:32:26 +0000329/* The teedataobject pre-allocates space for LINKCELLS number of objects.
330 To help the object fit neatly inside cache lines (space for 16 to 32
331 pointers), the value should be a multiple of 16 minus space for
332 the other structure members including PyHEAD overhead. The larger the
333 value, the less memory overhead per object and the less time spent
334 allocating/deallocating new links. The smaller the number, the less
335 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000337#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000338
339typedef struct {
340 PyObject_HEAD
341 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342 int numread;
343 PyObject *nextlink;
344 PyObject *(values[LINKCELLS]);
345} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000346
347typedef struct {
348 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000349 teedataobject *dataobj;
350 int index;
351} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000352
Raymond Hettingerad983e72003-11-12 14:32:26 +0000353static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000354
355static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000356teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000357{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000358 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000359
Raymond Hettingerad983e72003-11-12 14:32:26 +0000360 tdo = PyObject_New(teedataobject, &teedataobject_type);
361 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000362 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000363
364 tdo->numread = 0;
365 tdo->nextlink = NULL;
366 Py_INCREF(it);
367 tdo->it = it;
368 return (PyObject *)tdo;
369}
370
371static PyObject *
372teedataobject_jumplink(teedataobject *tdo)
373{
374 if (tdo->nextlink == NULL)
375 tdo->nextlink = teedataobject_new(tdo->it);
376 Py_INCREF(tdo->nextlink);
377 return tdo->nextlink;
378}
379
380static PyObject *
381teedataobject_getitem(teedataobject *tdo, int i)
382{
383 PyObject *value;
384
385 assert(i < LINKCELLS);
386 if (i < tdo->numread)
387 value = tdo->values[i];
388 else {
389 /* this is the lead iterator, so fetch more data */
390 assert(i == tdo->numread);
391 value = PyIter_Next(tdo->it);
392 if (value == NULL)
393 return NULL;
394 tdo->numread++;
395 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000396 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000397 Py_INCREF(value);
398 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000399}
400
401static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000402teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000403{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000404 int i;
Raymond Hettinger45143692003-10-25 06:37:47 +0000405
Raymond Hettingerad983e72003-11-12 14:32:26 +0000406 for (i=0 ; i<tdo->numread ; i++)
407 Py_DECREF(tdo->values[i]);
408 Py_XDECREF(tdo->it);
409 Py_XDECREF(tdo->nextlink);
410 PyObject_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000411}
412
Raymond Hettingerad983e72003-11-12 14:32:26 +0000413PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000414
Raymond Hettingerad983e72003-11-12 14:32:26 +0000415static PyTypeObject teedataobject_type = {
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000416 PyObject_HEAD_INIT(&PyType_Type)
417 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000418 "itertools.tee_dataobject", /* tp_name */
419 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000420 0, /* tp_itemsize */
421 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000422 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000423 0, /* tp_print */
424 0, /* tp_getattr */
425 0, /* tp_setattr */
426 0, /* tp_compare */
427 0, /* tp_repr */
428 0, /* tp_as_number */
429 0, /* tp_as_sequence */
430 0, /* tp_as_mapping */
431 0, /* tp_hash */
432 0, /* tp_call */
433 0, /* tp_str */
434 PyObject_GenericGetAttr, /* tp_getattro */
435 0, /* tp_setattro */
436 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000437 Py_TPFLAGS_DEFAULT, /* tp_flags */
438 teedataobject_doc, /* tp_doc */
439 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000440};
441
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000442
443static PyTypeObject tee_type;
444
445static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000446tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000447{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448 PyObject *value, *link;
449
450 if (to->index >= LINKCELLS) {
451 link = teedataobject_jumplink(to->dataobj);
452 Py_XDECREF(to->dataobj);
453 to->dataobj = (teedataobject *)link;
454 to->index = 0;
455 }
456 value = teedataobject_getitem(to->dataobj, to->index);
457 if (value == NULL)
458 return NULL;
459 to->index++;
460 return value;
461}
462
463static PyObject *
464tee_copy(teeobject *to)
465{
466 teeobject *newto;
467
468 newto = PyObject_New(teeobject, &tee_type);
469 if (newto == NULL)
470 return NULL;
471 Py_INCREF(to->dataobj);
472 newto->dataobj = to->dataobj;
473 newto->index = to->index;
474 return (PyObject *)newto;
475}
476
477PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
478
479static PyObject *
480tee_fromiterable(PyObject *iterable)
481{
482 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000483 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000484
485 it = PyObject_GetIter(iterable);
486 if (it == NULL)
487 return NULL;
488 if (PyObject_TypeCheck(it, &tee_type)) {
489 to = (teeobject *)tee_copy((teeobject *)it);
490 goto done;
491 }
492
493 to = PyObject_New(teeobject, &tee_type);
494 if (to == NULL)
495 goto done;
496 to->dataobj = (teedataobject *)teedataobject_new(it);
497 to->index = 0;
498done:
499 Py_XDECREF(it);
500 return (PyObject *)to;
501}
502
503static PyObject *
504tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
505{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000506 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000507
508 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
509 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000510 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000511}
512
513static void
514tee_dealloc(teeobject *to)
515{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000516 Py_XDECREF(to->dataobj);
517 PyObject_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000518}
519
Raymond Hettingerad983e72003-11-12 14:32:26 +0000520PyDoc_STRVAR(teeobject_doc,
521"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000522
Raymond Hettingerad983e72003-11-12 14:32:26 +0000523static PyMethodDef tee_methods[] = {
524 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
525 {NULL, NULL} /* sentinel */
526};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000527
528static PyTypeObject tee_type = {
529 PyObject_HEAD_INIT(NULL)
530 0, /* ob_size */
531 "itertools.tee", /* tp_name */
532 sizeof(teeobject), /* tp_basicsize */
533 0, /* tp_itemsize */
534 /* methods */
535 (destructor)tee_dealloc, /* tp_dealloc */
536 0, /* tp_print */
537 0, /* tp_getattr */
538 0, /* tp_setattr */
539 0, /* tp_compare */
540 0, /* tp_repr */
541 0, /* tp_as_number */
542 0, /* tp_as_sequence */
543 0, /* tp_as_mapping */
544 0, /* tp_hash */
545 0, /* tp_call */
546 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000547 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000548 0, /* tp_setattro */
549 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000550 Py_TPFLAGS_DEFAULT, /* tp_flags */
551 teeobject_doc, /* tp_doc */
552 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000553 0, /* tp_clear */
554 0, /* tp_richcompare */
555 0, /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000556 PyObject_SelfIter, /* tp_iter */
557 (iternextfunc)tee_next, /* tp_iternext */
558 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000559 0, /* tp_members */
560 0, /* tp_getset */
561 0, /* tp_base */
562 0, /* tp_dict */
563 0, /* tp_descr_get */
564 0, /* tp_descr_set */
565 0, /* tp_dictoffset */
566 0, /* tp_init */
567 0, /* tp_alloc */
568 tee_new, /* tp_new */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000569 PyObject_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000570};
571
Raymond Hettingerad983e72003-11-12 14:32:26 +0000572static PyObject *
573tee(PyObject *self, PyObject *args)
574{
575 int i, n=2;
576 PyObject *it, *iterable, *copyable, *result;
577
578 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
579 return NULL;
580 result = PyTuple_New(n);
581 if (result == NULL)
582 return NULL;
583 if (n == 0)
584 return result;
585 it = PyObject_GetIter(iterable);
586 if (it == NULL) {
587 Py_DECREF(result);
588 return NULL;
589 }
590 if (!PyObject_HasAttrString(it, "__copy__")) {
591 copyable = tee_fromiterable(it);
592 Py_DECREF(it);
593 if (copyable == NULL) {
594 Py_DECREF(result);
595 return NULL;
596 }
597 } else
598 copyable = it;
599 PyTuple_SET_ITEM(result, 0, copyable);
600 for (i=1 ; i<n ; i++) {
601 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
602 if (copyable == NULL) {
603 Py_DECREF(result);
604 return NULL;
605 }
606 PyTuple_SET_ITEM(result, i, copyable);
607 }
608 return result;
609}
610
611PyDoc_STRVAR(tee_doc,
612"tee(iterable, n=2) --> tuple of n independent iterators.");
613
614
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000615/* cycle object **********************************************************/
616
617typedef struct {
618 PyObject_HEAD
619 PyObject *it;
620 PyObject *saved;
621 int firstpass;
622} cycleobject;
623
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000624static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000625
626static PyObject *
627cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
628{
629 PyObject *it;
630 PyObject *iterable;
631 PyObject *saved;
632 cycleobject *lz;
633
634 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
635 return NULL;
636
637 /* Get iterator. */
638 it = PyObject_GetIter(iterable);
639 if (it == NULL)
640 return NULL;
641
642 saved = PyList_New(0);
643 if (saved == NULL) {
644 Py_DECREF(it);
645 return NULL;
646 }
647
648 /* create cycleobject structure */
649 lz = (cycleobject *)type->tp_alloc(type, 0);
650 if (lz == NULL) {
651 Py_DECREF(it);
652 Py_DECREF(saved);
653 return NULL;
654 }
655 lz->it = it;
656 lz->saved = saved;
657 lz->firstpass = 0;
658
659 return (PyObject *)lz;
660}
661
662static void
663cycle_dealloc(cycleobject *lz)
664{
665 PyObject_GC_UnTrack(lz);
666 Py_XDECREF(lz->saved);
667 Py_XDECREF(lz->it);
668 lz->ob_type->tp_free(lz);
669}
670
671static int
672cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
673{
674 int err;
675
676 if (lz->it) {
677 err = visit(lz->it, arg);
678 if (err)
679 return err;
680 }
681 if (lz->saved) {
682 err = visit(lz->saved, arg);
683 if (err)
684 return err;
685 }
686 return 0;
687}
688
689static PyObject *
690cycle_next(cycleobject *lz)
691{
692 PyObject *item;
693 PyObject *it;
694
695 while (1) {
696 item = PyIter_Next(lz->it);
697 if (item != NULL) {
698 if (!lz->firstpass)
699 PyList_Append(lz->saved, item);
700 return item;
701 }
702 if (PyList_Size(lz->saved) == 0)
703 return NULL;
704 it = PyObject_GetIter(lz->saved);
705 if (it == NULL)
706 return NULL;
707 Py_DECREF(lz->it);
708 lz->it = it;
709 lz->firstpass = 1;
710 }
711}
712
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000713PyDoc_STRVAR(cycle_doc,
714"cycle(iterable) --> cycle object\n\
715\n\
716Return elements from the iterable until it is exhausted.\n\
717Then repeat the sequence indefinitely.");
718
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000719static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000720 PyObject_HEAD_INIT(NULL)
721 0, /* ob_size */
722 "itertools.cycle", /* tp_name */
723 sizeof(cycleobject), /* tp_basicsize */
724 0, /* tp_itemsize */
725 /* methods */
726 (destructor)cycle_dealloc, /* tp_dealloc */
727 0, /* tp_print */
728 0, /* tp_getattr */
729 0, /* tp_setattr */
730 0, /* tp_compare */
731 0, /* tp_repr */
732 0, /* tp_as_number */
733 0, /* tp_as_sequence */
734 0, /* tp_as_mapping */
735 0, /* tp_hash */
736 0, /* tp_call */
737 0, /* tp_str */
738 PyObject_GenericGetAttr, /* tp_getattro */
739 0, /* tp_setattro */
740 0, /* tp_as_buffer */
741 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
742 Py_TPFLAGS_BASETYPE, /* tp_flags */
743 cycle_doc, /* tp_doc */
744 (traverseproc)cycle_traverse, /* tp_traverse */
745 0, /* tp_clear */
746 0, /* tp_richcompare */
747 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000748 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000749 (iternextfunc)cycle_next, /* tp_iternext */
750 0, /* tp_methods */
751 0, /* tp_members */
752 0, /* tp_getset */
753 0, /* tp_base */
754 0, /* tp_dict */
755 0, /* tp_descr_get */
756 0, /* tp_descr_set */
757 0, /* tp_dictoffset */
758 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000759 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760 cycle_new, /* tp_new */
761 PyObject_GC_Del, /* tp_free */
762};
763
764
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000765/* dropwhile object **********************************************************/
766
767typedef struct {
768 PyObject_HEAD
769 PyObject *func;
770 PyObject *it;
771 long start;
772} dropwhileobject;
773
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000774static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000775
776static PyObject *
777dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
778{
779 PyObject *func, *seq;
780 PyObject *it;
781 dropwhileobject *lz;
782
783 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
784 return NULL;
785
786 /* Get iterator. */
787 it = PyObject_GetIter(seq);
788 if (it == NULL)
789 return NULL;
790
791 /* create dropwhileobject structure */
792 lz = (dropwhileobject *)type->tp_alloc(type, 0);
793 if (lz == NULL) {
794 Py_DECREF(it);
795 return NULL;
796 }
797 Py_INCREF(func);
798 lz->func = func;
799 lz->it = it;
800 lz->start = 0;
801
802 return (PyObject *)lz;
803}
804
805static void
806dropwhile_dealloc(dropwhileobject *lz)
807{
808 PyObject_GC_UnTrack(lz);
809 Py_XDECREF(lz->func);
810 Py_XDECREF(lz->it);
811 lz->ob_type->tp_free(lz);
812}
813
814static int
815dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
816{
817 int err;
818
819 if (lz->it) {
820 err = visit(lz->it, arg);
821 if (err)
822 return err;
823 }
824 if (lz->func) {
825 err = visit(lz->func, arg);
826 if (err)
827 return err;
828 }
829 return 0;
830}
831
832static PyObject *
833dropwhile_next(dropwhileobject *lz)
834{
835 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000836 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000837 long ok;
838
839 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000840 assert(PyIter_Check(it));
841 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000842 if (item == NULL)
843 return NULL;
844 if (lz->start == 1)
845 return item;
846
847 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
848 if (good == NULL) {
849 Py_DECREF(item);
850 return NULL;
851 }
852 ok = PyObject_IsTrue(good);
853 Py_DECREF(good);
854 if (!ok) {
855 lz->start = 1;
856 return item;
857 }
858 Py_DECREF(item);
859 }
860}
861
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000862PyDoc_STRVAR(dropwhile_doc,
863"dropwhile(predicate, iterable) --> dropwhile object\n\
864\n\
865Drop items from the iterable while predicate(item) is true.\n\
866Afterwards, return every element until the iterable is exhausted.");
867
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000868static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000869 PyObject_HEAD_INIT(NULL)
870 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000871 "itertools.dropwhile", /* tp_name */
872 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000873 0, /* tp_itemsize */
874 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000875 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876 0, /* tp_print */
877 0, /* tp_getattr */
878 0, /* tp_setattr */
879 0, /* tp_compare */
880 0, /* tp_repr */
881 0, /* tp_as_number */
882 0, /* tp_as_sequence */
883 0, /* tp_as_mapping */
884 0, /* tp_hash */
885 0, /* tp_call */
886 0, /* tp_str */
887 PyObject_GenericGetAttr, /* tp_getattro */
888 0, /* tp_setattro */
889 0, /* tp_as_buffer */
890 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
891 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000892 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000893 (traverseproc)dropwhile_traverse, /* tp_traverse */
894 0, /* tp_clear */
895 0, /* tp_richcompare */
896 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000897 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000898 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000899 0, /* tp_methods */
900 0, /* tp_members */
901 0, /* tp_getset */
902 0, /* tp_base */
903 0, /* tp_dict */
904 0, /* tp_descr_get */
905 0, /* tp_descr_set */
906 0, /* tp_dictoffset */
907 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000908 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000909 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000910 PyObject_GC_Del, /* tp_free */
911};
912
913
914/* takewhile object **********************************************************/
915
916typedef struct {
917 PyObject_HEAD
918 PyObject *func;
919 PyObject *it;
920 long stop;
921} takewhileobject;
922
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000923static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000924
925static PyObject *
926takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
927{
928 PyObject *func, *seq;
929 PyObject *it;
930 takewhileobject *lz;
931
932 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
933 return NULL;
934
935 /* Get iterator. */
936 it = PyObject_GetIter(seq);
937 if (it == NULL)
938 return NULL;
939
940 /* create takewhileobject structure */
941 lz = (takewhileobject *)type->tp_alloc(type, 0);
942 if (lz == NULL) {
943 Py_DECREF(it);
944 return NULL;
945 }
946 Py_INCREF(func);
947 lz->func = func;
948 lz->it = it;
949 lz->stop = 0;
950
951 return (PyObject *)lz;
952}
953
954static void
955takewhile_dealloc(takewhileobject *lz)
956{
957 PyObject_GC_UnTrack(lz);
958 Py_XDECREF(lz->func);
959 Py_XDECREF(lz->it);
960 lz->ob_type->tp_free(lz);
961}
962
963static int
964takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
965{
966 int err;
967
968 if (lz->it) {
969 err = visit(lz->it, arg);
970 if (err)
971 return err;
972 }
973 if (lz->func) {
974 err = visit(lz->func, arg);
975 if (err)
976 return err;
977 }
978 return 0;
979}
980
981static PyObject *
982takewhile_next(takewhileobject *lz)
983{
984 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000985 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000986 long ok;
987
988 if (lz->stop == 1)
989 return NULL;
990
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000991 assert(PyIter_Check(it));
992 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000993 if (item == NULL)
994 return NULL;
995
996 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
997 if (good == NULL) {
998 Py_DECREF(item);
999 return NULL;
1000 }
1001 ok = PyObject_IsTrue(good);
1002 Py_DECREF(good);
1003 if (ok)
1004 return item;
1005 Py_DECREF(item);
1006 lz->stop = 1;
1007 return NULL;
1008}
1009
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001010PyDoc_STRVAR(takewhile_doc,
1011"takewhile(predicate, iterable) --> takewhile object\n\
1012\n\
1013Return successive entries from an iterable as long as the \n\
1014predicate evaluates to true for each entry.");
1015
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001016static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001017 PyObject_HEAD_INIT(NULL)
1018 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001019 "itertools.takewhile", /* tp_name */
1020 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001021 0, /* tp_itemsize */
1022 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001023 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001024 0, /* tp_print */
1025 0, /* tp_getattr */
1026 0, /* tp_setattr */
1027 0, /* tp_compare */
1028 0, /* tp_repr */
1029 0, /* tp_as_number */
1030 0, /* tp_as_sequence */
1031 0, /* tp_as_mapping */
1032 0, /* tp_hash */
1033 0, /* tp_call */
1034 0, /* tp_str */
1035 PyObject_GenericGetAttr, /* tp_getattro */
1036 0, /* tp_setattro */
1037 0, /* tp_as_buffer */
1038 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1039 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001040 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001041 (traverseproc)takewhile_traverse, /* tp_traverse */
1042 0, /* tp_clear */
1043 0, /* tp_richcompare */
1044 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001045 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001046 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001047 0, /* tp_methods */
1048 0, /* tp_members */
1049 0, /* tp_getset */
1050 0, /* tp_base */
1051 0, /* tp_dict */
1052 0, /* tp_descr_get */
1053 0, /* tp_descr_set */
1054 0, /* tp_dictoffset */
1055 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001056 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001057 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001058 PyObject_GC_Del, /* tp_free */
1059};
1060
1061
1062/* islice object ************************************************************/
1063
1064typedef struct {
1065 PyObject_HEAD
1066 PyObject *it;
1067 long next;
1068 long stop;
1069 long step;
1070 long cnt;
1071} isliceobject;
1072
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001073static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001074
1075static PyObject *
1076islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1077{
1078 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001079 long start=0, stop=-1, step=1;
1080 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001081 int numargs;
1082 isliceobject *lz;
1083
1084 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +00001085 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001086 return NULL;
1087
1088 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001089 if (a1 != Py_None) {
1090 stop = PyInt_AsLong(a1);
1091 if (stop == -1) {
1092 if (PyErr_Occurred())
1093 PyErr_Clear();
1094 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001095 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001096 return NULL;
1097 }
1098 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001099 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001100 start = PyInt_AsLong(a1);
1101 if (start == -1 && PyErr_Occurred()) {
1102 PyErr_Clear();
1103 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001104 "Start argument must be a non-negative integer.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001105 return NULL;
1106 }
1107 if (a2 != Py_None) {
1108 stop = PyInt_AsLong(a2);
1109 if (stop == -1) {
1110 if (PyErr_Occurred())
1111 PyErr_Clear();
1112 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001113 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001114 return NULL;
1115 }
1116 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001117 }
1118
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001119 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001120 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001121 "Indices for islice() must be non-negative integers.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001122 return NULL;
1123 }
1124
1125 if (step<1) {
1126 PyErr_SetString(PyExc_ValueError,
1127 "Step must be one or larger for islice().");
1128 return NULL;
1129 }
1130
1131 /* Get iterator. */
1132 it = PyObject_GetIter(seq);
1133 if (it == NULL)
1134 return NULL;
1135
1136 /* create isliceobject structure */
1137 lz = (isliceobject *)type->tp_alloc(type, 0);
1138 if (lz == NULL) {
1139 Py_DECREF(it);
1140 return NULL;
1141 }
1142 lz->it = it;
1143 lz->next = start;
1144 lz->stop = stop;
1145 lz->step = step;
1146 lz->cnt = 0L;
1147
1148 return (PyObject *)lz;
1149}
1150
1151static void
1152islice_dealloc(isliceobject *lz)
1153{
1154 PyObject_GC_UnTrack(lz);
1155 Py_XDECREF(lz->it);
1156 lz->ob_type->tp_free(lz);
1157}
1158
1159static int
1160islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1161{
1162 if (lz->it)
1163 return visit(lz->it, arg);
1164 return 0;
1165}
1166
1167static PyObject *
1168islice_next(isliceobject *lz)
1169{
1170 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001171 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001172 long oldnext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001173
1174 while (lz->cnt < lz->next) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001175 assert(PyIter_Check(it));
1176 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177 if (item == NULL)
1178 return NULL;
1179 Py_DECREF(item);
1180 lz->cnt++;
1181 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001182 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001183 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001184 assert(PyIter_Check(it));
1185 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001186 if (item == NULL)
1187 return NULL;
1188 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001189 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001190 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001191 if (lz->next < oldnext) /* Check for overflow */
1192 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001193 return item;
1194}
1195
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196PyDoc_STRVAR(islice_doc,
1197"islice(iterable, [start,] stop [, step]) --> islice object\n\
1198\n\
1199Return an iterator whose next() method returns selected values from an\n\
1200iterable. If start is specified, will skip all preceding elements;\n\
1201otherwise, start defaults to zero. Step defaults to one. If\n\
1202specified as another value, step determines how many values are \n\
1203skipped between successive calls. Works like a slice() on a list\n\
1204but returns an iterator.");
1205
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001206static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001207 PyObject_HEAD_INIT(NULL)
1208 0, /* ob_size */
1209 "itertools.islice", /* tp_name */
1210 sizeof(isliceobject), /* tp_basicsize */
1211 0, /* tp_itemsize */
1212 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001213 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001214 0, /* tp_print */
1215 0, /* tp_getattr */
1216 0, /* tp_setattr */
1217 0, /* tp_compare */
1218 0, /* tp_repr */
1219 0, /* tp_as_number */
1220 0, /* tp_as_sequence */
1221 0, /* tp_as_mapping */
1222 0, /* tp_hash */
1223 0, /* tp_call */
1224 0, /* tp_str */
1225 PyObject_GenericGetAttr, /* tp_getattro */
1226 0, /* tp_setattro */
1227 0, /* tp_as_buffer */
1228 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1229 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001230 islice_doc, /* tp_doc */
1231 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001232 0, /* tp_clear */
1233 0, /* tp_richcompare */
1234 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001235 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001236 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237 0, /* tp_methods */
1238 0, /* tp_members */
1239 0, /* tp_getset */
1240 0, /* tp_base */
1241 0, /* tp_dict */
1242 0, /* tp_descr_get */
1243 0, /* tp_descr_set */
1244 0, /* tp_dictoffset */
1245 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001246 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001247 islice_new, /* tp_new */
1248 PyObject_GC_Del, /* tp_free */
1249};
1250
1251
1252/* starmap object ************************************************************/
1253
1254typedef struct {
1255 PyObject_HEAD
1256 PyObject *func;
1257 PyObject *it;
1258} starmapobject;
1259
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001260static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001261
1262static PyObject *
1263starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1264{
1265 PyObject *func, *seq;
1266 PyObject *it;
1267 starmapobject *lz;
1268
1269 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1270 return NULL;
1271
1272 /* Get iterator. */
1273 it = PyObject_GetIter(seq);
1274 if (it == NULL)
1275 return NULL;
1276
1277 /* create starmapobject structure */
1278 lz = (starmapobject *)type->tp_alloc(type, 0);
1279 if (lz == NULL) {
1280 Py_DECREF(it);
1281 return NULL;
1282 }
1283 Py_INCREF(func);
1284 lz->func = func;
1285 lz->it = it;
1286
1287 return (PyObject *)lz;
1288}
1289
1290static void
1291starmap_dealloc(starmapobject *lz)
1292{
1293 PyObject_GC_UnTrack(lz);
1294 Py_XDECREF(lz->func);
1295 Py_XDECREF(lz->it);
1296 lz->ob_type->tp_free(lz);
1297}
1298
1299static int
1300starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1301{
1302 int err;
1303
1304 if (lz->it) {
1305 err = visit(lz->it, arg);
1306 if (err)
1307 return err;
1308 }
1309 if (lz->func) {
1310 err = visit(lz->func, arg);
1311 if (err)
1312 return err;
1313 }
1314 return 0;
1315}
1316
1317static PyObject *
1318starmap_next(starmapobject *lz)
1319{
1320 PyObject *args;
1321 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001322 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001323
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001324 assert(PyIter_Check(it));
1325 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001326 if (args == NULL)
1327 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001328 if (!PyTuple_CheckExact(args)) {
1329 Py_DECREF(args);
1330 PyErr_SetString(PyExc_TypeError,
1331 "iterator must return a tuple");
1332 return NULL;
1333 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334 result = PyObject_Call(lz->func, args, NULL);
1335 Py_DECREF(args);
1336 return result;
1337}
1338
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339PyDoc_STRVAR(starmap_doc,
1340"starmap(function, sequence) --> starmap object\n\
1341\n\
1342Return an iterator whose values are returned from the function evaluated\n\
1343with a argument tuple taken from the given sequence.");
1344
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001345static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346 PyObject_HEAD_INIT(NULL)
1347 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001348 "itertools.starmap", /* tp_name */
1349 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350 0, /* tp_itemsize */
1351 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001352 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353 0, /* tp_print */
1354 0, /* tp_getattr */
1355 0, /* tp_setattr */
1356 0, /* tp_compare */
1357 0, /* tp_repr */
1358 0, /* tp_as_number */
1359 0, /* tp_as_sequence */
1360 0, /* tp_as_mapping */
1361 0, /* tp_hash */
1362 0, /* tp_call */
1363 0, /* tp_str */
1364 PyObject_GenericGetAttr, /* tp_getattro */
1365 0, /* tp_setattro */
1366 0, /* tp_as_buffer */
1367 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1368 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001369 starmap_doc, /* tp_doc */
1370 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001371 0, /* tp_clear */
1372 0, /* tp_richcompare */
1373 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001374 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001375 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376 0, /* tp_methods */
1377 0, /* tp_members */
1378 0, /* tp_getset */
1379 0, /* tp_base */
1380 0, /* tp_dict */
1381 0, /* tp_descr_get */
1382 0, /* tp_descr_set */
1383 0, /* tp_dictoffset */
1384 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001385 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001386 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387 PyObject_GC_Del, /* tp_free */
1388};
1389
1390
1391/* imap object ************************************************************/
1392
1393typedef struct {
1394 PyObject_HEAD
1395 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001396 PyObject *func;
1397} imapobject;
1398
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001399static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001400
1401static PyObject *
1402imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1403{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001404 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001405 imapobject *lz;
1406 int numargs, i;
1407
1408 numargs = PyTuple_Size(args);
1409 if (numargs < 2) {
1410 PyErr_SetString(PyExc_TypeError,
1411 "imap() must have at least two arguments.");
1412 return NULL;
1413 }
1414
1415 iters = PyTuple_New(numargs-1);
1416 if (iters == NULL)
1417 return NULL;
1418
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001419 for (i=1 ; i<numargs ; i++) {
1420 /* Get iterator. */
1421 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1422 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423 Py_DECREF(iters);
1424 return NULL;
1425 }
1426 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001427 }
1428
1429 /* create imapobject structure */
1430 lz = (imapobject *)type->tp_alloc(type, 0);
1431 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001432 Py_DECREF(iters);
1433 return NULL;
1434 }
1435 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436 func = PyTuple_GET_ITEM(args, 0);
1437 Py_INCREF(func);
1438 lz->func = func;
1439
1440 return (PyObject *)lz;
1441}
1442
1443static void
1444imap_dealloc(imapobject *lz)
1445{
1446 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001447 Py_XDECREF(lz->iters);
1448 Py_XDECREF(lz->func);
1449 lz->ob_type->tp_free(lz);
1450}
1451
1452static int
1453imap_traverse(imapobject *lz, visitproc visit, void *arg)
1454{
1455 int err;
1456
1457 if (lz->iters) {
1458 err = visit(lz->iters, arg);
1459 if (err)
1460 return err;
1461 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462 if (lz->func) {
1463 err = visit(lz->func, arg);
1464 if (err)
1465 return err;
1466 }
1467 return 0;
1468}
1469
Raymond Hettinger2012f172003-02-07 05:32:58 +00001470/*
1471imap() is an iterator version of __builtins__.map() except that it does
1472not have the None fill-in feature. That was intentionally left out for
1473the following reasons:
1474
1475 1) Itertools are designed to be easily combined and chained together.
1476 Having all tools stop with the shortest input is a unifying principle
1477 that makes it easier to combine finite iterators (supplying data) with
1478 infinite iterators like count() and repeat() (for supplying sequential
1479 or constant arguments to a function).
1480
1481 2) In typical use cases for combining itertools, having one finite data
1482 supplier run out before another is likely to be an error condition which
1483 should not pass silently by automatically supplying None.
1484
1485 3) The use cases for automatic None fill-in are rare -- not many functions
1486 do something useful when a parameter suddenly switches type and becomes
1487 None.
1488
1489 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001490 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001491
1492 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1493*/
1494
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001495static PyObject *
1496imap_next(imapobject *lz)
1497{
1498 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001499 PyObject *argtuple;
1500 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001501 int numargs, i;
1502
1503 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001504 argtuple = PyTuple_New(numargs);
1505 if (argtuple == NULL)
1506 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001507
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001508 for (i=0 ; i<numargs ; i++) {
1509 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1510 if (val == NULL) {
1511 Py_DECREF(argtuple);
1512 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001513 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001514 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001515 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001516 if (lz->func == Py_None)
1517 return argtuple;
1518 result = PyObject_Call(lz->func, argtuple, NULL);
1519 Py_DECREF(argtuple);
1520 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001521}
1522
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523PyDoc_STRVAR(imap_doc,
1524"imap(func, *iterables) --> imap object\n\
1525\n\
1526Make an iterator that computes the function using arguments from\n\
1527each of the iterables. Like map() except that it returns\n\
1528an iterator instead of a list and that it stops when the shortest\n\
1529iterable is exhausted instead of filling in None for shorter\n\
1530iterables.");
1531
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001532static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533 PyObject_HEAD_INIT(NULL)
1534 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001535 "itertools.imap", /* tp_name */
1536 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537 0, /* tp_itemsize */
1538 /* methods */
1539 (destructor)imap_dealloc, /* tp_dealloc */
1540 0, /* tp_print */
1541 0, /* tp_getattr */
1542 0, /* tp_setattr */
1543 0, /* tp_compare */
1544 0, /* tp_repr */
1545 0, /* tp_as_number */
1546 0, /* tp_as_sequence */
1547 0, /* tp_as_mapping */
1548 0, /* tp_hash */
1549 0, /* tp_call */
1550 0, /* tp_str */
1551 PyObject_GenericGetAttr, /* tp_getattro */
1552 0, /* tp_setattro */
1553 0, /* tp_as_buffer */
1554 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1555 Py_TPFLAGS_BASETYPE, /* tp_flags */
1556 imap_doc, /* tp_doc */
1557 (traverseproc)imap_traverse, /* tp_traverse */
1558 0, /* tp_clear */
1559 0, /* tp_richcompare */
1560 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001561 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562 (iternextfunc)imap_next, /* tp_iternext */
1563 0, /* tp_methods */
1564 0, /* tp_members */
1565 0, /* tp_getset */
1566 0, /* tp_base */
1567 0, /* tp_dict */
1568 0, /* tp_descr_get */
1569 0, /* tp_descr_set */
1570 0, /* tp_dictoffset */
1571 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001572 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001573 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001574 PyObject_GC_Del, /* tp_free */
1575};
1576
1577
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001578/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001579
1580typedef struct {
1581 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001582 long tuplesize;
1583 long iternum; /* which iterator is active */
1584 PyObject *ittuple; /* tuple of iterators */
1585} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001586
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001587static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001588
1589static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001590chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001591{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001592 chainobject *lz;
1593 int tuplesize = PySequence_Length(args);
1594 int i;
1595 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001597 /* obtain iterators */
1598 assert(PyTuple_Check(args));
1599 ittuple = PyTuple_New(tuplesize);
1600 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001601 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001602 for (i=0; i < tuplesize; ++i) {
1603 PyObject *item = PyTuple_GET_ITEM(args, i);
1604 PyObject *it = PyObject_GetIter(item);
1605 if (it == NULL) {
1606 if (PyErr_ExceptionMatches(PyExc_TypeError))
1607 PyErr_Format(PyExc_TypeError,
1608 "chain argument #%d must support iteration",
1609 i+1);
1610 Py_DECREF(ittuple);
1611 return NULL;
1612 }
1613 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001614 }
1615
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001616 /* create chainobject structure */
1617 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001618 if (lz == NULL) {
1619 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001620 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001621 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001622
1623 lz->ittuple = ittuple;
1624 lz->iternum = 0;
1625 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001626
1627 return (PyObject *)lz;
1628}
1629
1630static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001631chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001632{
1633 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001634 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001635 lz->ob_type->tp_free(lz);
1636}
1637
Raymond Hettinger2012f172003-02-07 05:32:58 +00001638static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001639chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001640{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001641 if (lz->ittuple)
1642 return visit(lz->ittuple, arg);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001643 return 0;
1644}
1645
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001646static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001647chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001648{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001649 PyObject *it;
1650 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001651
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001652 while (lz->iternum < lz->tuplesize) {
1653 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1654 item = PyIter_Next(it);
1655 if (item != NULL)
1656 return item;
1657 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001658 }
1659 return NULL;
1660}
1661
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001662PyDoc_STRVAR(chain_doc,
1663"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001664\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001665Return a chain object whose .next() method returns elements from the\n\
1666first iterable until it is exhausted, then elements from the next\n\
1667iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001668
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001669static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670 PyObject_HEAD_INIT(NULL)
1671 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001672 "itertools.chain", /* tp_name */
1673 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674 0, /* tp_itemsize */
1675 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001676 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001677 0, /* tp_print */
1678 0, /* tp_getattr */
1679 0, /* tp_setattr */
1680 0, /* tp_compare */
1681 0, /* tp_repr */
1682 0, /* tp_as_number */
1683 0, /* tp_as_sequence */
1684 0, /* tp_as_mapping */
1685 0, /* tp_hash */
1686 0, /* tp_call */
1687 0, /* tp_str */
1688 PyObject_GenericGetAttr, /* tp_getattro */
1689 0, /* tp_setattro */
1690 0, /* tp_as_buffer */
1691 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1692 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001693 chain_doc, /* tp_doc */
1694 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001695 0, /* tp_clear */
1696 0, /* tp_richcompare */
1697 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001698 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001699 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001700 0, /* tp_methods */
1701 0, /* tp_members */
1702 0, /* tp_getset */
1703 0, /* tp_base */
1704 0, /* tp_dict */
1705 0, /* tp_descr_get */
1706 0, /* tp_descr_set */
1707 0, /* tp_dictoffset */
1708 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001709 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001710 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001711 PyObject_GC_Del, /* tp_free */
1712};
1713
1714
1715/* ifilter object ************************************************************/
1716
1717typedef struct {
1718 PyObject_HEAD
1719 PyObject *func;
1720 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001721} ifilterobject;
1722
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001723static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724
1725static PyObject *
1726ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1727{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001728 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001729 PyObject *it;
1730 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001731
Raymond Hettinger60eca932003-02-09 06:40:58 +00001732 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001733 return NULL;
1734
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735 /* Get iterator. */
1736 it = PyObject_GetIter(seq);
1737 if (it == NULL)
1738 return NULL;
1739
1740 /* create ifilterobject structure */
1741 lz = (ifilterobject *)type->tp_alloc(type, 0);
1742 if (lz == NULL) {
1743 Py_DECREF(it);
1744 return NULL;
1745 }
1746 Py_INCREF(func);
1747 lz->func = func;
1748 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749
1750 return (PyObject *)lz;
1751}
1752
1753static void
1754ifilter_dealloc(ifilterobject *lz)
1755{
1756 PyObject_GC_UnTrack(lz);
1757 Py_XDECREF(lz->func);
1758 Py_XDECREF(lz->it);
1759 lz->ob_type->tp_free(lz);
1760}
1761
1762static int
1763ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1764{
1765 int err;
1766
1767 if (lz->it) {
1768 err = visit(lz->it, arg);
1769 if (err)
1770 return err;
1771 }
1772 if (lz->func) {
1773 err = visit(lz->func, arg);
1774 if (err)
1775 return err;
1776 }
1777 return 0;
1778}
1779
1780static PyObject *
1781ifilter_next(ifilterobject *lz)
1782{
1783 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001784 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001785 long ok;
1786
1787 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001788 assert(PyIter_Check(it));
1789 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001790 if (item == NULL)
1791 return NULL;
1792
1793 if (lz->func == Py_None) {
1794 ok = PyObject_IsTrue(item);
1795 } else {
1796 PyObject *good;
1797 good = PyObject_CallFunctionObjArgs(lz->func,
1798 item, NULL);
1799 if (good == NULL) {
1800 Py_DECREF(item);
1801 return NULL;
1802 }
1803 ok = PyObject_IsTrue(good);
1804 Py_DECREF(good);
1805 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001806 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001807 return item;
1808 Py_DECREF(item);
1809 }
1810}
1811
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001812PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001813"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001814\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001815Return those items of sequence for which function(item) is true.\n\
1816If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001817
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001818static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001819 PyObject_HEAD_INIT(NULL)
1820 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001821 "itertools.ifilter", /* tp_name */
1822 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001823 0, /* tp_itemsize */
1824 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001825 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001826 0, /* tp_print */
1827 0, /* tp_getattr */
1828 0, /* tp_setattr */
1829 0, /* tp_compare */
1830 0, /* tp_repr */
1831 0, /* tp_as_number */
1832 0, /* tp_as_sequence */
1833 0, /* tp_as_mapping */
1834 0, /* tp_hash */
1835 0, /* tp_call */
1836 0, /* tp_str */
1837 PyObject_GenericGetAttr, /* tp_getattro */
1838 0, /* tp_setattro */
1839 0, /* tp_as_buffer */
1840 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1841 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001842 ifilter_doc, /* tp_doc */
1843 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001844 0, /* tp_clear */
1845 0, /* tp_richcompare */
1846 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001847 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001848 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001849 0, /* tp_methods */
1850 0, /* tp_members */
1851 0, /* tp_getset */
1852 0, /* tp_base */
1853 0, /* tp_dict */
1854 0, /* tp_descr_get */
1855 0, /* tp_descr_set */
1856 0, /* tp_dictoffset */
1857 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001858 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001859 ifilter_new, /* tp_new */
1860 PyObject_GC_Del, /* tp_free */
1861};
1862
1863
1864/* ifilterfalse object ************************************************************/
1865
1866typedef struct {
1867 PyObject_HEAD
1868 PyObject *func;
1869 PyObject *it;
1870} ifilterfalseobject;
1871
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001872static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001873
1874static PyObject *
1875ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1876{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001877 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001878 PyObject *it;
1879 ifilterfalseobject *lz;
1880
1881 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1882 return NULL;
1883
1884 /* Get iterator. */
1885 it = PyObject_GetIter(seq);
1886 if (it == NULL)
1887 return NULL;
1888
1889 /* create ifilterfalseobject structure */
1890 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1891 if (lz == NULL) {
1892 Py_DECREF(it);
1893 return NULL;
1894 }
1895 Py_INCREF(func);
1896 lz->func = func;
1897 lz->it = it;
1898
1899 return (PyObject *)lz;
1900}
1901
1902static void
1903ifilterfalse_dealloc(ifilterfalseobject *lz)
1904{
1905 PyObject_GC_UnTrack(lz);
1906 Py_XDECREF(lz->func);
1907 Py_XDECREF(lz->it);
1908 lz->ob_type->tp_free(lz);
1909}
1910
1911static int
1912ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1913{
1914 int err;
1915
1916 if (lz->it) {
1917 err = visit(lz->it, arg);
1918 if (err)
1919 return err;
1920 }
1921 if (lz->func) {
1922 err = visit(lz->func, arg);
1923 if (err)
1924 return err;
1925 }
1926 return 0;
1927}
1928
1929static PyObject *
1930ifilterfalse_next(ifilterfalseobject *lz)
1931{
1932 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001933 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001934 long ok;
1935
1936 for (;;) {
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001937 assert(PyIter_Check(it));
1938 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001939 if (item == NULL)
1940 return NULL;
1941
1942 if (lz->func == Py_None) {
1943 ok = PyObject_IsTrue(item);
1944 } else {
1945 PyObject *good;
1946 good = PyObject_CallFunctionObjArgs(lz->func,
1947 item, NULL);
1948 if (good == NULL) {
1949 Py_DECREF(item);
1950 return NULL;
1951 }
1952 ok = PyObject_IsTrue(good);
1953 Py_DECREF(good);
1954 }
1955 if (!ok)
1956 return item;
1957 Py_DECREF(item);
1958 }
1959}
1960
Raymond Hettinger60eca932003-02-09 06:40:58 +00001961PyDoc_STRVAR(ifilterfalse_doc,
1962"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1963\n\
1964Return those items of sequence for which function(item) is false.\n\
1965If function is None, return the items that are false.");
1966
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001967static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001968 PyObject_HEAD_INIT(NULL)
1969 0, /* ob_size */
1970 "itertools.ifilterfalse", /* tp_name */
1971 sizeof(ifilterfalseobject), /* tp_basicsize */
1972 0, /* tp_itemsize */
1973 /* methods */
1974 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1975 0, /* tp_print */
1976 0, /* tp_getattr */
1977 0, /* tp_setattr */
1978 0, /* tp_compare */
1979 0, /* tp_repr */
1980 0, /* tp_as_number */
1981 0, /* tp_as_sequence */
1982 0, /* tp_as_mapping */
1983 0, /* tp_hash */
1984 0, /* tp_call */
1985 0, /* tp_str */
1986 PyObject_GenericGetAttr, /* tp_getattro */
1987 0, /* tp_setattro */
1988 0, /* tp_as_buffer */
1989 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1990 Py_TPFLAGS_BASETYPE, /* tp_flags */
1991 ifilterfalse_doc, /* tp_doc */
1992 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1993 0, /* tp_clear */
1994 0, /* tp_richcompare */
1995 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001996 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001997 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1998 0, /* tp_methods */
1999 0, /* tp_members */
2000 0, /* tp_getset */
2001 0, /* tp_base */
2002 0, /* tp_dict */
2003 0, /* tp_descr_get */
2004 0, /* tp_descr_set */
2005 0, /* tp_dictoffset */
2006 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002007 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002008 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002009 PyObject_GC_Del, /* tp_free */
2010};
2011
2012
2013/* count object ************************************************************/
2014
2015typedef struct {
2016 PyObject_HEAD
2017 long cnt;
2018} countobject;
2019
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002020static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002021
2022static PyObject *
2023count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2024{
2025 countobject *lz;
2026 long cnt = 0;
2027
2028 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
2029 return NULL;
2030
2031 /* create countobject structure */
2032 lz = (countobject *)PyObject_New(countobject, &count_type);
2033 if (lz == NULL)
2034 return NULL;
2035 lz->cnt = cnt;
2036
2037 return (PyObject *)lz;
2038}
2039
2040static PyObject *
2041count_next(countobject *lz)
2042{
2043 return PyInt_FromLong(lz->cnt++);
2044}
2045
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002046PyDoc_STRVAR(count_doc,
2047"count([firstval]) --> count object\n\
2048\n\
2049Return a count object whose .next() method returns consecutive\n\
2050integers starting from zero or, if specified, from firstval.");
2051
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002052static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002053 PyObject_HEAD_INIT(NULL)
2054 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002055 "itertools.count", /* tp_name */
2056 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002057 0, /* tp_itemsize */
2058 /* methods */
2059 (destructor)PyObject_Del, /* tp_dealloc */
2060 0, /* tp_print */
2061 0, /* tp_getattr */
2062 0, /* tp_setattr */
2063 0, /* tp_compare */
2064 0, /* tp_repr */
2065 0, /* tp_as_number */
2066 0, /* tp_as_sequence */
2067 0, /* tp_as_mapping */
2068 0, /* tp_hash */
2069 0, /* tp_call */
2070 0, /* tp_str */
2071 PyObject_GenericGetAttr, /* tp_getattro */
2072 0, /* tp_setattro */
2073 0, /* tp_as_buffer */
2074 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002075 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002076 0, /* tp_traverse */
2077 0, /* tp_clear */
2078 0, /* tp_richcompare */
2079 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002080 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002081 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002082 0, /* tp_methods */
2083 0, /* tp_members */
2084 0, /* tp_getset */
2085 0, /* tp_base */
2086 0, /* tp_dict */
2087 0, /* tp_descr_get */
2088 0, /* tp_descr_set */
2089 0, /* tp_dictoffset */
2090 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002091 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002092 count_new, /* tp_new */
2093};
2094
2095
2096/* izip object ************************************************************/
2097
2098#include "Python.h"
2099
2100typedef struct {
2101 PyObject_HEAD
2102 long tuplesize;
2103 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002104 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002105} izipobject;
2106
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002107static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002108
2109static PyObject *
2110izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2111{
2112 izipobject *lz;
2113 int i;
2114 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002115 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002116 int tuplesize = PySequence_Length(args);
2117
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002118 /* args must be a tuple */
2119 assert(PyTuple_Check(args));
2120
2121 /* obtain iterators */
2122 ittuple = PyTuple_New(tuplesize);
2123 if(ittuple == NULL)
2124 return NULL;
2125 for (i=0; i < tuplesize; ++i) {
2126 PyObject *item = PyTuple_GET_ITEM(args, i);
2127 PyObject *it = PyObject_GetIter(item);
2128 if (it == NULL) {
2129 if (PyErr_ExceptionMatches(PyExc_TypeError))
2130 PyErr_Format(PyExc_TypeError,
2131 "izip argument #%d must support iteration",
2132 i+1);
2133 Py_DECREF(ittuple);
2134 return NULL;
2135 }
2136 PyTuple_SET_ITEM(ittuple, i, it);
2137 }
2138
Raymond Hettinger2012f172003-02-07 05:32:58 +00002139 /* create a result holder */
2140 result = PyTuple_New(tuplesize);
2141 if (result == NULL) {
2142 Py_DECREF(ittuple);
2143 return NULL;
2144 }
2145 for (i=0 ; i < tuplesize ; i++) {
2146 Py_INCREF(Py_None);
2147 PyTuple_SET_ITEM(result, i, Py_None);
2148 }
2149
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002150 /* create izipobject structure */
2151 lz = (izipobject *)type->tp_alloc(type, 0);
2152 if (lz == NULL) {
2153 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002154 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002155 return NULL;
2156 }
2157 lz->ittuple = ittuple;
2158 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002159 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002160
2161 return (PyObject *)lz;
2162}
2163
2164static void
2165izip_dealloc(izipobject *lz)
2166{
2167 PyObject_GC_UnTrack(lz);
2168 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002169 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002170 lz->ob_type->tp_free(lz);
2171}
2172
2173static int
2174izip_traverse(izipobject *lz, visitproc visit, void *arg)
2175{
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00002176 int err;
2177
2178 if (lz->ittuple) {
2179 err = visit(lz->ittuple, arg);
2180 if (err)
2181 return err;
2182 }
2183 if (lz->result) {
2184 err = visit(lz->result, arg);
2185 if (err)
2186 return err;
2187 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002188 return 0;
2189}
2190
2191static PyObject *
2192izip_next(izipobject *lz)
2193{
2194 int i;
2195 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002196 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002197 PyObject *it;
2198 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002199 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002200
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002201 if (tuplesize == 0)
2202 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002203 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002204 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002205 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002206 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002207 assert(PyIter_Check(it));
2208 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002209 if (item == NULL) {
2210 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002211 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002212 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002213 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002214 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002215 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002216 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002217 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002218 result = PyTuple_New(tuplesize);
2219 if (result == NULL)
2220 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002221 for (i=0 ; i < tuplesize ; i++) {
2222 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002223 assert(PyIter_Check(it));
2224 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002225 if (item == NULL) {
2226 Py_DECREF(result);
2227 return NULL;
2228 }
2229 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002230 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002231 }
2232 return result;
2233}
2234
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002235PyDoc_STRVAR(izip_doc,
2236"izip(iter1 [,iter2 [...]]) --> izip object\n\
2237\n\
2238Return a izip object whose .next() method returns a tuple where\n\
2239the i-th element comes from the i-th iterable argument. The .next()\n\
2240method continues until the shortest iterable in the argument sequence\n\
2241is exhausted and then it raises StopIteration. Works like the zip()\n\
2242function but consumes less memory by returning an iterator instead of\n\
2243a list.");
2244
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002245static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002246 PyObject_HEAD_INIT(NULL)
2247 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002248 "itertools.izip", /* tp_name */
2249 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002250 0, /* tp_itemsize */
2251 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002252 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002253 0, /* tp_print */
2254 0, /* tp_getattr */
2255 0, /* tp_setattr */
2256 0, /* tp_compare */
2257 0, /* tp_repr */
2258 0, /* tp_as_number */
2259 0, /* tp_as_sequence */
2260 0, /* tp_as_mapping */
2261 0, /* tp_hash */
2262 0, /* tp_call */
2263 0, /* tp_str */
2264 PyObject_GenericGetAttr, /* tp_getattro */
2265 0, /* tp_setattro */
2266 0, /* tp_as_buffer */
2267 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2268 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002269 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002270 (traverseproc)izip_traverse, /* tp_traverse */
2271 0, /* tp_clear */
2272 0, /* tp_richcompare */
2273 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002274 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002275 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002276 0, /* tp_methods */
2277 0, /* tp_members */
2278 0, /* tp_getset */
2279 0, /* tp_base */
2280 0, /* tp_dict */
2281 0, /* tp_descr_get */
2282 0, /* tp_descr_set */
2283 0, /* tp_dictoffset */
2284 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002285 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002286 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002287 PyObject_GC_Del, /* tp_free */
2288};
2289
2290
2291/* repeat object ************************************************************/
2292
2293typedef struct {
2294 PyObject_HEAD
2295 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002296 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002297} repeatobject;
2298
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002299static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002300
2301static PyObject *
2302repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2303{
2304 repeatobject *ro;
2305 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002306 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002307
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002308 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002309 return NULL;
2310
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002311 if (PyTuple_Size(args) == 2 && cnt < 0)
2312 cnt = 0;
2313
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002314 ro = (repeatobject *)type->tp_alloc(type, 0);
2315 if (ro == NULL)
2316 return NULL;
2317 Py_INCREF(element);
2318 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002319 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002320 return (PyObject *)ro;
2321}
2322
2323static void
2324repeat_dealloc(repeatobject *ro)
2325{
2326 PyObject_GC_UnTrack(ro);
2327 Py_XDECREF(ro->element);
2328 ro->ob_type->tp_free(ro);
2329}
2330
2331static int
2332repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2333{
2334 if (ro->element)
2335 return visit(ro->element, arg);
2336 return 0;
2337}
2338
2339static PyObject *
2340repeat_next(repeatobject *ro)
2341{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002342 if (ro->cnt == 0)
2343 return NULL;
2344 if (ro->cnt > 0)
2345 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002346 Py_INCREF(ro->element);
2347 return ro->element;
2348}
2349
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002350PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002351"repeat(element [,times]) -> create an iterator which returns the element\n\
2352for the specified number of times. If not specified, returns the element\n\
2353endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002354
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002355static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002356 PyObject_HEAD_INIT(NULL)
2357 0, /* ob_size */
2358 "itertools.repeat", /* tp_name */
2359 sizeof(repeatobject), /* tp_basicsize */
2360 0, /* tp_itemsize */
2361 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002362 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002363 0, /* tp_print */
2364 0, /* tp_getattr */
2365 0, /* tp_setattr */
2366 0, /* tp_compare */
2367 0, /* tp_repr */
2368 0, /* tp_as_number */
2369 0, /* tp_as_sequence */
2370 0, /* tp_as_mapping */
2371 0, /* tp_hash */
2372 0, /* tp_call */
2373 0, /* tp_str */
2374 PyObject_GenericGetAttr, /* tp_getattro */
2375 0, /* tp_setattro */
2376 0, /* tp_as_buffer */
2377 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2378 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002379 repeat_doc, /* tp_doc */
2380 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002381 0, /* tp_clear */
2382 0, /* tp_richcompare */
2383 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002384 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002385 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002386 0, /* tp_methods */
2387 0, /* tp_members */
2388 0, /* tp_getset */
2389 0, /* tp_base */
2390 0, /* tp_dict */
2391 0, /* tp_descr_get */
2392 0, /* tp_descr_set */
2393 0, /* tp_dictoffset */
2394 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002395 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002396 repeat_new, /* tp_new */
2397 PyObject_GC_Del, /* tp_free */
2398};
2399
2400
2401/* module level code ********************************************************/
2402
2403PyDoc_STRVAR(module_doc,
2404"Functional tools for creating and using iterators.\n\
2405\n\
2406Infinite iterators:\n\
2407count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002408cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002409repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002410\n\
2411Iterators terminating on the shortest input sequence:\n\
2412izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002413ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2414ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002415islice(seq, [start,] stop [, step]) --> elements from\n\
2416 seq[start:stop:step]\n\
2417imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2418starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002419tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002420chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002421takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2422dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002423groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002424");
2425
2426
Raymond Hettingerad983e72003-11-12 14:32:26 +00002427static PyMethodDef module_methods[] = {
2428 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2429 {NULL, NULL} /* sentinel */
2430};
2431
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002432PyMODINIT_FUNC
2433inititertools(void)
2434{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002435 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002436 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002437 char *name;
2438 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002439 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002440 &dropwhile_type,
2441 &takewhile_type,
2442 &islice_type,
2443 &starmap_type,
2444 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002445 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002446 &ifilter_type,
2447 &ifilterfalse_type,
2448 &count_type,
2449 &izip_type,
2450 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002451 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002452 NULL
2453 };
2454
Raymond Hettingerad983e72003-11-12 14:32:26 +00002455 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002456
Raymond Hettinger60eca932003-02-09 06:40:58 +00002457 for (i=0 ; typelist[i] != NULL ; i++) {
2458 if (PyType_Ready(typelist[i]) < 0)
2459 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002460 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002461 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002462 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002463 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002464 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002465
2466 if (PyType_Ready(&teedataobject_type) < 0)
2467 return;
2468 if (PyType_Ready(&tee_type) < 0)
2469 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002470 if (PyType_Ready(&_grouper_type) < 0)
2471 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002472}