blob: bf2c49311ffb36154f59e0ab40e020f1342f53ef [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{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +000067 Py_VISIT(gbo->it);
68 Py_VISIT(gbo->keyfunc);
69 Py_VISIT(gbo->tgtkey);
70 Py_VISIT(gbo->currkey);
71 Py_VISIT(gbo->currvalue);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000072 return 0;
73}
74
75static PyObject *
76groupby_next(groupbyobject *gbo)
77{
78 PyObject *newvalue, *newkey, *r, *grouper;
79
80 /* skip to next iteration group */
81 for (;;) {
82 if (gbo->currkey == NULL)
83 /* pass */;
84 else if (gbo->tgtkey == NULL)
85 break;
86 else {
87 int rcmp;
88
89 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
90 gbo->currkey, Py_EQ);
91 if (rcmp == -1)
92 return NULL;
93 else if (rcmp == 0)
94 break;
95 }
96
97 newvalue = PyIter_Next(gbo->it);
98 if (newvalue == NULL)
99 return NULL;
100
101 if (gbo->keyfunc == Py_None) {
102 newkey = newvalue;
103 Py_INCREF(newvalue);
104 } else {
105 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
106 newvalue, NULL);
107 if (newkey == NULL) {
108 Py_DECREF(newvalue);
109 return NULL;
110 }
111 }
112
113 Py_XDECREF(gbo->currkey);
114 gbo->currkey = newkey;
115 Py_XDECREF(gbo->currvalue);
116 gbo->currvalue = newvalue;
117 }
118
119 Py_XDECREF(gbo->tgtkey);
120 gbo->tgtkey = gbo->currkey;
121 Py_INCREF(gbo->currkey);
122
123 grouper = _grouper_create(gbo, gbo->tgtkey);
124 if (grouper == NULL)
125 return NULL;
126
127 r = PyTuple_Pack(2, gbo->currkey, grouper);
128 Py_DECREF(grouper);
129 return r;
130}
131
132PyDoc_STRVAR(groupby_doc,
133"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
134(key, sub-iterator) grouped by each value of key(value).\n");
135
136static PyTypeObject groupby_type = {
137 PyObject_HEAD_INIT(NULL)
138 0, /* ob_size */
139 "itertools.groupby", /* tp_name */
140 sizeof(groupbyobject), /* tp_basicsize */
141 0, /* tp_itemsize */
142 /* methods */
143 (destructor)groupby_dealloc, /* tp_dealloc */
144 0, /* tp_print */
145 0, /* tp_getattr */
146 0, /* tp_setattr */
147 0, /* tp_compare */
148 0, /* tp_repr */
149 0, /* tp_as_number */
150 0, /* tp_as_sequence */
151 0, /* tp_as_mapping */
152 0, /* tp_hash */
153 0, /* tp_call */
154 0, /* tp_str */
155 PyObject_GenericGetAttr, /* tp_getattro */
156 0, /* tp_setattro */
157 0, /* tp_as_buffer */
158 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
159 Py_TPFLAGS_BASETYPE, /* tp_flags */
160 groupby_doc, /* tp_doc */
161 (traverseproc)groupby_traverse, /* tp_traverse */
162 0, /* tp_clear */
163 0, /* tp_richcompare */
164 0, /* tp_weaklistoffset */
165 PyObject_SelfIter, /* tp_iter */
166 (iternextfunc)groupby_next, /* tp_iternext */
167 0, /* tp_methods */
168 0, /* tp_members */
169 0, /* tp_getset */
170 0, /* tp_base */
171 0, /* tp_dict */
172 0, /* tp_descr_get */
173 0, /* tp_descr_set */
174 0, /* tp_dictoffset */
175 0, /* tp_init */
176 0, /* tp_alloc */
177 groupby_new, /* tp_new */
178 PyObject_GC_Del, /* tp_free */
179};
180
181
182/* _grouper object (internal) ************************************************/
183
184typedef struct {
185 PyObject_HEAD
186 PyObject *parent;
187 PyObject *tgtkey;
188} _grouperobject;
189
190static PyTypeObject _grouper_type;
191
192static PyObject *
193_grouper_create(groupbyobject *parent, PyObject *tgtkey)
194{
195 _grouperobject *igo;
196
197 igo = PyObject_New(_grouperobject, &_grouper_type);
198 if (igo == NULL)
199 return NULL;
200 igo->parent = (PyObject *)parent;
201 Py_INCREF(parent);
202 igo->tgtkey = tgtkey;
203 Py_INCREF(tgtkey);
204
205 return (PyObject *)igo;
206}
207
208static void
209_grouper_dealloc(_grouperobject *igo)
210{
211 Py_DECREF(igo->parent);
212 Py_DECREF(igo->tgtkey);
213 PyObject_Del(igo);
214}
215
216static PyObject *
217_grouper_next(_grouperobject *igo)
218{
219 groupbyobject *gbo = (groupbyobject *)igo->parent;
220 PyObject *newvalue, *newkey, *r;
221 int rcmp;
222
223 if (gbo->currvalue == NULL) {
224 newvalue = PyIter_Next(gbo->it);
225 if (newvalue == NULL)
226 return NULL;
227
228 if (gbo->keyfunc == Py_None) {
229 newkey = newvalue;
230 Py_INCREF(newvalue);
231 } else {
232 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
233 newvalue, NULL);
234 if (newkey == NULL) {
235 Py_DECREF(newvalue);
236 return NULL;
237 }
238 }
239
240 assert(gbo->currkey == NULL);
241 gbo->currkey = newkey;
242 gbo->currvalue = newvalue;
243 }
244
245 assert(gbo->currkey != NULL);
246 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
247 if (rcmp <= 0)
248 /* got any error or current group is end */
249 return NULL;
250
251 r = gbo->currvalue;
252 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000253 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000254
255 return r;
256}
257
258static PyTypeObject _grouper_type = {
259 PyObject_HEAD_INIT(NULL)
260 0, /* ob_size */
261 "itertools._grouper", /* tp_name */
262 sizeof(_grouperobject), /* tp_basicsize */
263 0, /* tp_itemsize */
264 /* methods */
265 (destructor)_grouper_dealloc, /* tp_dealloc */
266 0, /* tp_print */
267 0, /* tp_getattr */
268 0, /* tp_setattr */
269 0, /* tp_compare */
270 0, /* tp_repr */
271 0, /* tp_as_number */
272 0, /* tp_as_sequence */
273 0, /* tp_as_mapping */
274 0, /* tp_hash */
275 0, /* tp_call */
276 0, /* tp_str */
277 PyObject_GenericGetAttr, /* tp_getattro */
278 0, /* tp_setattro */
279 0, /* tp_as_buffer */
280 Py_TPFLAGS_DEFAULT, /* tp_flags */
281 0, /* tp_doc */
282 0, /* tp_traverse */
283 0, /* tp_clear */
284 0, /* tp_richcompare */
285 0, /* tp_weaklistoffset */
286 PyObject_SelfIter, /* tp_iter */
287 (iternextfunc)_grouper_next, /* tp_iternext */
288 0, /* tp_methods */
289 0, /* tp_members */
290 0, /* tp_getset */
291 0, /* tp_base */
292 0, /* tp_dict */
293 0, /* tp_descr_get */
294 0, /* tp_descr_set */
295 0, /* tp_dictoffset */
296 0, /* tp_init */
297 0, /* tp_alloc */
298 0, /* tp_new */
299 PyObject_Del, /* tp_free */
300};
301
302
303
Raymond Hettingerad983e72003-11-12 14:32:26 +0000304/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000305
Raymond Hettingerad983e72003-11-12 14:32:26 +0000306/* The teedataobject pre-allocates space for LINKCELLS number of objects.
307 To help the object fit neatly inside cache lines (space for 16 to 32
308 pointers), the value should be a multiple of 16 minus space for
309 the other structure members including PyHEAD overhead. The larger the
310 value, the less memory overhead per object and the less time spent
311 allocating/deallocating new links. The smaller the number, the less
312 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000313*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000314#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000315
316typedef struct {
317 PyObject_HEAD
318 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000319 int numread;
320 PyObject *nextlink;
321 PyObject *(values[LINKCELLS]);
322} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000323
324typedef struct {
325 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000326 teedataobject *dataobj;
327 int index;
328} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000329
Raymond Hettingerad983e72003-11-12 14:32:26 +0000330static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000331
332static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000333teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000334{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000335 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000336
Raymond Hettingerad983e72003-11-12 14:32:26 +0000337 tdo = PyObject_New(teedataobject, &teedataobject_type);
338 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000339 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000340
341 tdo->numread = 0;
342 tdo->nextlink = NULL;
343 Py_INCREF(it);
344 tdo->it = it;
345 return (PyObject *)tdo;
346}
347
348static PyObject *
349teedataobject_jumplink(teedataobject *tdo)
350{
351 if (tdo->nextlink == NULL)
352 tdo->nextlink = teedataobject_new(tdo->it);
353 Py_INCREF(tdo->nextlink);
354 return tdo->nextlink;
355}
356
357static PyObject *
358teedataobject_getitem(teedataobject *tdo, int i)
359{
360 PyObject *value;
361
362 assert(i < LINKCELLS);
363 if (i < tdo->numread)
364 value = tdo->values[i];
365 else {
366 /* this is the lead iterator, so fetch more data */
367 assert(i == tdo->numread);
368 value = PyIter_Next(tdo->it);
369 if (value == NULL)
370 return NULL;
371 tdo->numread++;
372 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000373 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000374 Py_INCREF(value);
375 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000376}
377
378static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000379teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000380{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000381 int i;
Raymond Hettinger45143692003-10-25 06:37:47 +0000382
Raymond Hettingerad983e72003-11-12 14:32:26 +0000383 for (i=0 ; i<tdo->numread ; i++)
384 Py_DECREF(tdo->values[i]);
385 Py_XDECREF(tdo->it);
386 Py_XDECREF(tdo->nextlink);
387 PyObject_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000388}
389
Raymond Hettingerad983e72003-11-12 14:32:26 +0000390PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000391
Raymond Hettingerad983e72003-11-12 14:32:26 +0000392static PyTypeObject teedataobject_type = {
Skip Montanarof3938fd2004-02-10 20:27:40 +0000393 PyObject_HEAD_INIT(0) /* Must fill in type value later */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000394 0, /* ob_size */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000395 "itertools.tee_dataobject", /* tp_name */
396 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000397 0, /* tp_itemsize */
398 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000400 0, /* tp_print */
401 0, /* tp_getattr */
402 0, /* tp_setattr */
403 0, /* tp_compare */
404 0, /* tp_repr */
405 0, /* tp_as_number */
406 0, /* tp_as_sequence */
407 0, /* tp_as_mapping */
408 0, /* tp_hash */
409 0, /* tp_call */
410 0, /* tp_str */
411 PyObject_GenericGetAttr, /* tp_getattro */
412 0, /* tp_setattro */
413 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000414 Py_TPFLAGS_DEFAULT, /* tp_flags */
415 teedataobject_doc, /* tp_doc */
416 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000417};
418
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000419
420static PyTypeObject tee_type;
421
422static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000423tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000424{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000425 PyObject *value, *link;
426
427 if (to->index >= LINKCELLS) {
428 link = teedataobject_jumplink(to->dataobj);
429 Py_XDECREF(to->dataobj);
430 to->dataobj = (teedataobject *)link;
431 to->index = 0;
432 }
433 value = teedataobject_getitem(to->dataobj, to->index);
434 if (value == NULL)
435 return NULL;
436 to->index++;
437 return value;
438}
439
440static PyObject *
441tee_copy(teeobject *to)
442{
443 teeobject *newto;
444
445 newto = PyObject_New(teeobject, &tee_type);
446 if (newto == NULL)
447 return NULL;
448 Py_INCREF(to->dataobj);
449 newto->dataobj = to->dataobj;
450 newto->index = to->index;
451 return (PyObject *)newto;
452}
453
454PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
455
456static PyObject *
457tee_fromiterable(PyObject *iterable)
458{
459 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000460 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000461
462 it = PyObject_GetIter(iterable);
463 if (it == NULL)
464 return NULL;
465 if (PyObject_TypeCheck(it, &tee_type)) {
466 to = (teeobject *)tee_copy((teeobject *)it);
467 goto done;
468 }
469
470 to = PyObject_New(teeobject, &tee_type);
471 if (to == NULL)
472 goto done;
473 to->dataobj = (teedataobject *)teedataobject_new(it);
474 to->index = 0;
475done:
476 Py_XDECREF(it);
477 return (PyObject *)to;
478}
479
480static PyObject *
481tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
482{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000483 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000484
485 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
486 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000487 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000488}
489
490static void
491tee_dealloc(teeobject *to)
492{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000493 Py_XDECREF(to->dataobj);
494 PyObject_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000495}
496
Raymond Hettingerad983e72003-11-12 14:32:26 +0000497PyDoc_STRVAR(teeobject_doc,
498"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000499
Raymond Hettingerad983e72003-11-12 14:32:26 +0000500static PyMethodDef tee_methods[] = {
501 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
502 {NULL, NULL} /* sentinel */
503};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000504
505static PyTypeObject tee_type = {
506 PyObject_HEAD_INIT(NULL)
507 0, /* ob_size */
508 "itertools.tee", /* tp_name */
509 sizeof(teeobject), /* tp_basicsize */
510 0, /* tp_itemsize */
511 /* methods */
512 (destructor)tee_dealloc, /* tp_dealloc */
513 0, /* tp_print */
514 0, /* tp_getattr */
515 0, /* tp_setattr */
516 0, /* tp_compare */
517 0, /* tp_repr */
518 0, /* tp_as_number */
519 0, /* tp_as_sequence */
520 0, /* tp_as_mapping */
521 0, /* tp_hash */
522 0, /* tp_call */
523 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000524 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000525 0, /* tp_setattro */
526 0, /* tp_as_buffer */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000527 Py_TPFLAGS_DEFAULT, /* tp_flags */
528 teeobject_doc, /* tp_doc */
529 0, /* tp_traverse */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000530 0, /* tp_clear */
531 0, /* tp_richcompare */
532 0, /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000533 PyObject_SelfIter, /* tp_iter */
534 (iternextfunc)tee_next, /* tp_iternext */
535 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000536 0, /* tp_members */
537 0, /* tp_getset */
538 0, /* tp_base */
539 0, /* tp_dict */
540 0, /* tp_descr_get */
541 0, /* tp_descr_set */
542 0, /* tp_dictoffset */
543 0, /* tp_init */
544 0, /* tp_alloc */
545 tee_new, /* tp_new */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000546 PyObject_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000547};
548
Raymond Hettingerad983e72003-11-12 14:32:26 +0000549static PyObject *
550tee(PyObject *self, PyObject *args)
551{
552 int i, n=2;
553 PyObject *it, *iterable, *copyable, *result;
554
555 if (!PyArg_ParseTuple(args, "O|i", &iterable, &n))
556 return NULL;
557 result = PyTuple_New(n);
558 if (result == NULL)
559 return NULL;
560 if (n == 0)
561 return result;
562 it = PyObject_GetIter(iterable);
563 if (it == NULL) {
564 Py_DECREF(result);
565 return NULL;
566 }
567 if (!PyObject_HasAttrString(it, "__copy__")) {
568 copyable = tee_fromiterable(it);
569 Py_DECREF(it);
570 if (copyable == NULL) {
571 Py_DECREF(result);
572 return NULL;
573 }
574 } else
575 copyable = it;
576 PyTuple_SET_ITEM(result, 0, copyable);
577 for (i=1 ; i<n ; i++) {
578 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
579 if (copyable == NULL) {
580 Py_DECREF(result);
581 return NULL;
582 }
583 PyTuple_SET_ITEM(result, i, copyable);
584 }
585 return result;
586}
587
588PyDoc_STRVAR(tee_doc,
589"tee(iterable, n=2) --> tuple of n independent iterators.");
590
591
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000592/* cycle object **********************************************************/
593
594typedef struct {
595 PyObject_HEAD
596 PyObject *it;
597 PyObject *saved;
598 int firstpass;
599} cycleobject;
600
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000601static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000602
603static PyObject *
604cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
605{
606 PyObject *it;
607 PyObject *iterable;
608 PyObject *saved;
609 cycleobject *lz;
610
611 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
612 return NULL;
613
614 /* Get iterator. */
615 it = PyObject_GetIter(iterable);
616 if (it == NULL)
617 return NULL;
618
619 saved = PyList_New(0);
620 if (saved == NULL) {
621 Py_DECREF(it);
622 return NULL;
623 }
624
625 /* create cycleobject structure */
626 lz = (cycleobject *)type->tp_alloc(type, 0);
627 if (lz == NULL) {
628 Py_DECREF(it);
629 Py_DECREF(saved);
630 return NULL;
631 }
632 lz->it = it;
633 lz->saved = saved;
634 lz->firstpass = 0;
635
636 return (PyObject *)lz;
637}
638
639static void
640cycle_dealloc(cycleobject *lz)
641{
642 PyObject_GC_UnTrack(lz);
643 Py_XDECREF(lz->saved);
644 Py_XDECREF(lz->it);
645 lz->ob_type->tp_free(lz);
646}
647
648static int
649cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
650{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000651 Py_VISIT(lz->it);
652 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000653 return 0;
654}
655
656static PyObject *
657cycle_next(cycleobject *lz)
658{
659 PyObject *item;
660 PyObject *it;
661
662 while (1) {
663 item = PyIter_Next(lz->it);
664 if (item != NULL) {
665 if (!lz->firstpass)
666 PyList_Append(lz->saved, item);
667 return item;
668 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000669 if (PyErr_Occurred()) {
670 if (PyErr_ExceptionMatches(PyExc_StopIteration))
671 PyErr_Clear();
672 else
673 return NULL;
674 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000675 if (PyList_Size(lz->saved) == 0)
676 return NULL;
677 it = PyObject_GetIter(lz->saved);
678 if (it == NULL)
679 return NULL;
680 Py_DECREF(lz->it);
681 lz->it = it;
682 lz->firstpass = 1;
683 }
684}
685
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000686PyDoc_STRVAR(cycle_doc,
687"cycle(iterable) --> cycle object\n\
688\n\
689Return elements from the iterable until it is exhausted.\n\
690Then repeat the sequence indefinitely.");
691
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000692static PyTypeObject cycle_type = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000693 PyObject_HEAD_INIT(NULL)
694 0, /* ob_size */
695 "itertools.cycle", /* tp_name */
696 sizeof(cycleobject), /* tp_basicsize */
697 0, /* tp_itemsize */
698 /* methods */
699 (destructor)cycle_dealloc, /* tp_dealloc */
700 0, /* tp_print */
701 0, /* tp_getattr */
702 0, /* tp_setattr */
703 0, /* tp_compare */
704 0, /* tp_repr */
705 0, /* tp_as_number */
706 0, /* tp_as_sequence */
707 0, /* tp_as_mapping */
708 0, /* tp_hash */
709 0, /* tp_call */
710 0, /* tp_str */
711 PyObject_GenericGetAttr, /* tp_getattro */
712 0, /* tp_setattro */
713 0, /* tp_as_buffer */
714 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
715 Py_TPFLAGS_BASETYPE, /* tp_flags */
716 cycle_doc, /* tp_doc */
717 (traverseproc)cycle_traverse, /* tp_traverse */
718 0, /* tp_clear */
719 0, /* tp_richcompare */
720 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000721 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000722 (iternextfunc)cycle_next, /* tp_iternext */
723 0, /* tp_methods */
724 0, /* tp_members */
725 0, /* tp_getset */
726 0, /* tp_base */
727 0, /* tp_dict */
728 0, /* tp_descr_get */
729 0, /* tp_descr_set */
730 0, /* tp_dictoffset */
731 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000732 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000733 cycle_new, /* tp_new */
734 PyObject_GC_Del, /* tp_free */
735};
736
737
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000738/* dropwhile object **********************************************************/
739
740typedef struct {
741 PyObject_HEAD
742 PyObject *func;
743 PyObject *it;
744 long start;
745} dropwhileobject;
746
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000747static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000748
749static PyObject *
750dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
751{
752 PyObject *func, *seq;
753 PyObject *it;
754 dropwhileobject *lz;
755
756 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
757 return NULL;
758
759 /* Get iterator. */
760 it = PyObject_GetIter(seq);
761 if (it == NULL)
762 return NULL;
763
764 /* create dropwhileobject structure */
765 lz = (dropwhileobject *)type->tp_alloc(type, 0);
766 if (lz == NULL) {
767 Py_DECREF(it);
768 return NULL;
769 }
770 Py_INCREF(func);
771 lz->func = func;
772 lz->it = it;
773 lz->start = 0;
774
775 return (PyObject *)lz;
776}
777
778static void
779dropwhile_dealloc(dropwhileobject *lz)
780{
781 PyObject_GC_UnTrack(lz);
782 Py_XDECREF(lz->func);
783 Py_XDECREF(lz->it);
784 lz->ob_type->tp_free(lz);
785}
786
787static int
788dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
789{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000790 Py_VISIT(lz->it);
791 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000792 return 0;
793}
794
795static PyObject *
796dropwhile_next(dropwhileobject *lz)
797{
798 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000799 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000800 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000801 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000802
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000803 assert(PyIter_Check(it));
804 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000805 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000806 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000807 if (item == NULL)
808 return NULL;
809 if (lz->start == 1)
810 return item;
811
812 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
813 if (good == NULL) {
814 Py_DECREF(item);
815 return NULL;
816 }
817 ok = PyObject_IsTrue(good);
818 Py_DECREF(good);
819 if (!ok) {
820 lz->start = 1;
821 return item;
822 }
823 Py_DECREF(item);
824 }
825}
826
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000827PyDoc_STRVAR(dropwhile_doc,
828"dropwhile(predicate, iterable) --> dropwhile object\n\
829\n\
830Drop items from the iterable while predicate(item) is true.\n\
831Afterwards, return every element until the iterable is exhausted.");
832
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000833static PyTypeObject dropwhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000834 PyObject_HEAD_INIT(NULL)
835 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000836 "itertools.dropwhile", /* tp_name */
837 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000838 0, /* tp_itemsize */
839 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000840 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000841 0, /* tp_print */
842 0, /* tp_getattr */
843 0, /* tp_setattr */
844 0, /* tp_compare */
845 0, /* tp_repr */
846 0, /* tp_as_number */
847 0, /* tp_as_sequence */
848 0, /* tp_as_mapping */
849 0, /* tp_hash */
850 0, /* tp_call */
851 0, /* tp_str */
852 PyObject_GenericGetAttr, /* tp_getattro */
853 0, /* tp_setattro */
854 0, /* tp_as_buffer */
855 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
856 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000857 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000858 (traverseproc)dropwhile_traverse, /* tp_traverse */
859 0, /* tp_clear */
860 0, /* tp_richcompare */
861 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000862 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000863 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000864 0, /* tp_methods */
865 0, /* tp_members */
866 0, /* tp_getset */
867 0, /* tp_base */
868 0, /* tp_dict */
869 0, /* tp_descr_get */
870 0, /* tp_descr_set */
871 0, /* tp_dictoffset */
872 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000873 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000874 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000875 PyObject_GC_Del, /* tp_free */
876};
877
878
879/* takewhile object **********************************************************/
880
881typedef struct {
882 PyObject_HEAD
883 PyObject *func;
884 PyObject *it;
885 long stop;
886} takewhileobject;
887
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000888static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000889
890static PyObject *
891takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
892{
893 PyObject *func, *seq;
894 PyObject *it;
895 takewhileobject *lz;
896
897 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
898 return NULL;
899
900 /* Get iterator. */
901 it = PyObject_GetIter(seq);
902 if (it == NULL)
903 return NULL;
904
905 /* create takewhileobject structure */
906 lz = (takewhileobject *)type->tp_alloc(type, 0);
907 if (lz == NULL) {
908 Py_DECREF(it);
909 return NULL;
910 }
911 Py_INCREF(func);
912 lz->func = func;
913 lz->it = it;
914 lz->stop = 0;
915
916 return (PyObject *)lz;
917}
918
919static void
920takewhile_dealloc(takewhileobject *lz)
921{
922 PyObject_GC_UnTrack(lz);
923 Py_XDECREF(lz->func);
924 Py_XDECREF(lz->it);
925 lz->ob_type->tp_free(lz);
926}
927
928static int
929takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
930{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000931 Py_VISIT(lz->it);
932 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000933 return 0;
934}
935
936static PyObject *
937takewhile_next(takewhileobject *lz)
938{
939 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000940 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000941 long ok;
942
943 if (lz->stop == 1)
944 return NULL;
945
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000946 assert(PyIter_Check(it));
947 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000948 if (item == NULL)
949 return NULL;
950
951 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
952 if (good == NULL) {
953 Py_DECREF(item);
954 return NULL;
955 }
956 ok = PyObject_IsTrue(good);
957 Py_DECREF(good);
958 if (ok)
959 return item;
960 Py_DECREF(item);
961 lz->stop = 1;
962 return NULL;
963}
964
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000965PyDoc_STRVAR(takewhile_doc,
966"takewhile(predicate, iterable) --> takewhile object\n\
967\n\
968Return successive entries from an iterable as long as the \n\
969predicate evaluates to true for each entry.");
970
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000971static PyTypeObject takewhile_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000972 PyObject_HEAD_INIT(NULL)
973 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000974 "itertools.takewhile", /* tp_name */
975 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000976 0, /* tp_itemsize */
977 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000978 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000979 0, /* tp_print */
980 0, /* tp_getattr */
981 0, /* tp_setattr */
982 0, /* tp_compare */
983 0, /* tp_repr */
984 0, /* tp_as_number */
985 0, /* tp_as_sequence */
986 0, /* tp_as_mapping */
987 0, /* tp_hash */
988 0, /* tp_call */
989 0, /* tp_str */
990 PyObject_GenericGetAttr, /* tp_getattro */
991 0, /* tp_setattro */
992 0, /* tp_as_buffer */
993 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
994 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000995 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000996 (traverseproc)takewhile_traverse, /* tp_traverse */
997 0, /* tp_clear */
998 0, /* tp_richcompare */
999 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001000 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001001 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001002 0, /* tp_methods */
1003 0, /* tp_members */
1004 0, /* tp_getset */
1005 0, /* tp_base */
1006 0, /* tp_dict */
1007 0, /* tp_descr_get */
1008 0, /* tp_descr_set */
1009 0, /* tp_dictoffset */
1010 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001011 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001012 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001013 PyObject_GC_Del, /* tp_free */
1014};
1015
1016
1017/* islice object ************************************************************/
1018
1019typedef struct {
1020 PyObject_HEAD
1021 PyObject *it;
1022 long next;
1023 long stop;
1024 long step;
1025 long cnt;
1026} isliceobject;
1027
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001028static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001029
1030static PyObject *
1031islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1032{
1033 PyObject *seq;
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001034 long start=0, stop=-1, step=1;
1035 PyObject *it, *a1=NULL, *a2=NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001036 int numargs;
1037 isliceobject *lz;
1038
1039 numargs = PyTuple_Size(args);
Raymond Hettinger341deb72003-05-02 19:44:20 +00001040 if (!PyArg_ParseTuple(args, "OO|Ol:islice", &seq, &a1, &a2, &step))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001041 return NULL;
1042
1043 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001044 if (a1 != Py_None) {
1045 stop = PyInt_AsLong(a1);
1046 if (stop == -1) {
1047 if (PyErr_Occurred())
1048 PyErr_Clear();
1049 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001050 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001051 return NULL;
1052 }
1053 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001054 } else {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001055 start = PyInt_AsLong(a1);
1056 if (start == -1 && PyErr_Occurred()) {
1057 PyErr_Clear();
1058 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001059 "Start argument must be a non-negative integer.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001060 return NULL;
1061 }
1062 if (a2 != Py_None) {
1063 stop = PyInt_AsLong(a2);
1064 if (stop == -1) {
1065 if (PyErr_Occurred())
1066 PyErr_Clear();
1067 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001068 "Stop argument must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001069 return NULL;
1070 }
1071 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001072 }
1073
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001074 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001075 PyErr_SetString(PyExc_ValueError,
Raymond Hettinger0e4f7642003-10-28 07:32:28 +00001076 "Indices for islice() must be non-negative integers.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001077 return NULL;
1078 }
1079
1080 if (step<1) {
1081 PyErr_SetString(PyExc_ValueError,
1082 "Step must be one or larger for islice().");
1083 return NULL;
1084 }
1085
1086 /* Get iterator. */
1087 it = PyObject_GetIter(seq);
1088 if (it == NULL)
1089 return NULL;
1090
1091 /* create isliceobject structure */
1092 lz = (isliceobject *)type->tp_alloc(type, 0);
1093 if (lz == NULL) {
1094 Py_DECREF(it);
1095 return NULL;
1096 }
1097 lz->it = it;
1098 lz->next = start;
1099 lz->stop = stop;
1100 lz->step = step;
1101 lz->cnt = 0L;
1102
1103 return (PyObject *)lz;
1104}
1105
1106static void
1107islice_dealloc(isliceobject *lz)
1108{
1109 PyObject_GC_UnTrack(lz);
1110 Py_XDECREF(lz->it);
1111 lz->ob_type->tp_free(lz);
1112}
1113
1114static int
1115islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1116{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001117 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001118 return 0;
1119}
1120
1121static PyObject *
1122islice_next(isliceobject *lz)
1123{
1124 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001125 PyObject *it = lz->it;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001126 long oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001127 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001128
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001129 assert(PyIter_Check(it));
1130 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001131 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001132 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001133 if (item == NULL)
1134 return NULL;
1135 Py_DECREF(item);
1136 lz->cnt++;
1137 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001138 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001139 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001140 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001141 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001142 if (item == NULL)
1143 return NULL;
1144 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001145 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001146 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001147 if (lz->next < oldnext) /* Check for overflow */
1148 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001149 return item;
1150}
1151
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001152PyDoc_STRVAR(islice_doc,
1153"islice(iterable, [start,] stop [, step]) --> islice object\n\
1154\n\
1155Return an iterator whose next() method returns selected values from an\n\
1156iterable. If start is specified, will skip all preceding elements;\n\
1157otherwise, start defaults to zero. Step defaults to one. If\n\
1158specified as another value, step determines how many values are \n\
1159skipped between successive calls. Works like a slice() on a list\n\
1160but returns an iterator.");
1161
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001162static PyTypeObject islice_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163 PyObject_HEAD_INIT(NULL)
1164 0, /* ob_size */
1165 "itertools.islice", /* tp_name */
1166 sizeof(isliceobject), /* tp_basicsize */
1167 0, /* tp_itemsize */
1168 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001169 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001170 0, /* tp_print */
1171 0, /* tp_getattr */
1172 0, /* tp_setattr */
1173 0, /* tp_compare */
1174 0, /* tp_repr */
1175 0, /* tp_as_number */
1176 0, /* tp_as_sequence */
1177 0, /* tp_as_mapping */
1178 0, /* tp_hash */
1179 0, /* tp_call */
1180 0, /* tp_str */
1181 PyObject_GenericGetAttr, /* tp_getattro */
1182 0, /* tp_setattro */
1183 0, /* tp_as_buffer */
1184 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1185 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001186 islice_doc, /* tp_doc */
1187 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001188 0, /* tp_clear */
1189 0, /* tp_richcompare */
1190 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001191 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001192 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001193 0, /* tp_methods */
1194 0, /* tp_members */
1195 0, /* tp_getset */
1196 0, /* tp_base */
1197 0, /* tp_dict */
1198 0, /* tp_descr_get */
1199 0, /* tp_descr_set */
1200 0, /* tp_dictoffset */
1201 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001202 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001203 islice_new, /* tp_new */
1204 PyObject_GC_Del, /* tp_free */
1205};
1206
1207
1208/* starmap object ************************************************************/
1209
1210typedef struct {
1211 PyObject_HEAD
1212 PyObject *func;
1213 PyObject *it;
1214} starmapobject;
1215
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001216static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001217
1218static PyObject *
1219starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1220{
1221 PyObject *func, *seq;
1222 PyObject *it;
1223 starmapobject *lz;
1224
1225 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1226 return NULL;
1227
1228 /* Get iterator. */
1229 it = PyObject_GetIter(seq);
1230 if (it == NULL)
1231 return NULL;
1232
1233 /* create starmapobject structure */
1234 lz = (starmapobject *)type->tp_alloc(type, 0);
1235 if (lz == NULL) {
1236 Py_DECREF(it);
1237 return NULL;
1238 }
1239 Py_INCREF(func);
1240 lz->func = func;
1241 lz->it = it;
1242
1243 return (PyObject *)lz;
1244}
1245
1246static void
1247starmap_dealloc(starmapobject *lz)
1248{
1249 PyObject_GC_UnTrack(lz);
1250 Py_XDECREF(lz->func);
1251 Py_XDECREF(lz->it);
1252 lz->ob_type->tp_free(lz);
1253}
1254
1255static int
1256starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1257{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001258 Py_VISIT(lz->it);
1259 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001260 return 0;
1261}
1262
1263static PyObject *
1264starmap_next(starmapobject *lz)
1265{
1266 PyObject *args;
1267 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001268 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001270 assert(PyIter_Check(it));
1271 args = (*it->ob_type->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001272 if (args == NULL)
1273 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001274 if (!PyTuple_CheckExact(args)) {
1275 Py_DECREF(args);
1276 PyErr_SetString(PyExc_TypeError,
1277 "iterator must return a tuple");
1278 return NULL;
1279 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001280 result = PyObject_Call(lz->func, args, NULL);
1281 Py_DECREF(args);
1282 return result;
1283}
1284
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001285PyDoc_STRVAR(starmap_doc,
1286"starmap(function, sequence) --> starmap object\n\
1287\n\
1288Return an iterator whose values are returned from the function evaluated\n\
1289with a argument tuple taken from the given sequence.");
1290
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001291static PyTypeObject starmap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001292 PyObject_HEAD_INIT(NULL)
1293 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001294 "itertools.starmap", /* tp_name */
1295 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001296 0, /* tp_itemsize */
1297 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001298 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001299 0, /* tp_print */
1300 0, /* tp_getattr */
1301 0, /* tp_setattr */
1302 0, /* tp_compare */
1303 0, /* tp_repr */
1304 0, /* tp_as_number */
1305 0, /* tp_as_sequence */
1306 0, /* tp_as_mapping */
1307 0, /* tp_hash */
1308 0, /* tp_call */
1309 0, /* tp_str */
1310 PyObject_GenericGetAttr, /* tp_getattro */
1311 0, /* tp_setattro */
1312 0, /* tp_as_buffer */
1313 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1314 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001315 starmap_doc, /* tp_doc */
1316 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001317 0, /* tp_clear */
1318 0, /* tp_richcompare */
1319 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001320 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001321 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001322 0, /* tp_methods */
1323 0, /* tp_members */
1324 0, /* tp_getset */
1325 0, /* tp_base */
1326 0, /* tp_dict */
1327 0, /* tp_descr_get */
1328 0, /* tp_descr_set */
1329 0, /* tp_dictoffset */
1330 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001331 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001332 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001333 PyObject_GC_Del, /* tp_free */
1334};
1335
1336
1337/* imap object ************************************************************/
1338
1339typedef struct {
1340 PyObject_HEAD
1341 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342 PyObject *func;
1343} imapobject;
1344
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001345static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346
1347static PyObject *
1348imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1349{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001350 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351 imapobject *lz;
1352 int numargs, i;
1353
1354 numargs = PyTuple_Size(args);
1355 if (numargs < 2) {
1356 PyErr_SetString(PyExc_TypeError,
1357 "imap() must have at least two arguments.");
1358 return NULL;
1359 }
1360
1361 iters = PyTuple_New(numargs-1);
1362 if (iters == NULL)
1363 return NULL;
1364
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001365 for (i=1 ; i<numargs ; i++) {
1366 /* Get iterator. */
1367 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1368 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001369 Py_DECREF(iters);
1370 return NULL;
1371 }
1372 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001373 }
1374
1375 /* create imapobject structure */
1376 lz = (imapobject *)type->tp_alloc(type, 0);
1377 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001378 Py_DECREF(iters);
1379 return NULL;
1380 }
1381 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001382 func = PyTuple_GET_ITEM(args, 0);
1383 Py_INCREF(func);
1384 lz->func = func;
1385
1386 return (PyObject *)lz;
1387}
1388
1389static void
1390imap_dealloc(imapobject *lz)
1391{
1392 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001393 Py_XDECREF(lz->iters);
1394 Py_XDECREF(lz->func);
1395 lz->ob_type->tp_free(lz);
1396}
1397
1398static int
1399imap_traverse(imapobject *lz, visitproc visit, void *arg)
1400{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001401 Py_VISIT(lz->iters);
1402 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001403 return 0;
1404}
1405
Raymond Hettinger2012f172003-02-07 05:32:58 +00001406/*
1407imap() is an iterator version of __builtins__.map() except that it does
1408not have the None fill-in feature. That was intentionally left out for
1409the following reasons:
1410
1411 1) Itertools are designed to be easily combined and chained together.
1412 Having all tools stop with the shortest input is a unifying principle
1413 that makes it easier to combine finite iterators (supplying data) with
1414 infinite iterators like count() and repeat() (for supplying sequential
1415 or constant arguments to a function).
1416
1417 2) In typical use cases for combining itertools, having one finite data
1418 supplier run out before another is likely to be an error condition which
1419 should not pass silently by automatically supplying None.
1420
1421 3) The use cases for automatic None fill-in are rare -- not many functions
1422 do something useful when a parameter suddenly switches type and becomes
1423 None.
1424
1425 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001426 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001427
1428 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1429*/
1430
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001431static PyObject *
1432imap_next(imapobject *lz)
1433{
1434 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001435 PyObject *argtuple;
1436 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437 int numargs, i;
1438
1439 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001440 argtuple = PyTuple_New(numargs);
1441 if (argtuple == NULL)
1442 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001443
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001444 for (i=0 ; i<numargs ; i++) {
1445 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1446 if (val == NULL) {
1447 Py_DECREF(argtuple);
1448 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001449 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001450 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001451 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001452 if (lz->func == Py_None)
1453 return argtuple;
1454 result = PyObject_Call(lz->func, argtuple, NULL);
1455 Py_DECREF(argtuple);
1456 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457}
1458
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459PyDoc_STRVAR(imap_doc,
1460"imap(func, *iterables) --> imap object\n\
1461\n\
1462Make an iterator that computes the function using arguments from\n\
1463each of the iterables. Like map() except that it returns\n\
1464an iterator instead of a list and that it stops when the shortest\n\
1465iterable is exhausted instead of filling in None for shorter\n\
1466iterables.");
1467
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001468static PyTypeObject imap_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469 PyObject_HEAD_INIT(NULL)
1470 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001471 "itertools.imap", /* tp_name */
1472 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473 0, /* tp_itemsize */
1474 /* methods */
1475 (destructor)imap_dealloc, /* tp_dealloc */
1476 0, /* tp_print */
1477 0, /* tp_getattr */
1478 0, /* tp_setattr */
1479 0, /* tp_compare */
1480 0, /* tp_repr */
1481 0, /* tp_as_number */
1482 0, /* tp_as_sequence */
1483 0, /* tp_as_mapping */
1484 0, /* tp_hash */
1485 0, /* tp_call */
1486 0, /* tp_str */
1487 PyObject_GenericGetAttr, /* tp_getattro */
1488 0, /* tp_setattro */
1489 0, /* tp_as_buffer */
1490 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1491 Py_TPFLAGS_BASETYPE, /* tp_flags */
1492 imap_doc, /* tp_doc */
1493 (traverseproc)imap_traverse, /* tp_traverse */
1494 0, /* tp_clear */
1495 0, /* tp_richcompare */
1496 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001497 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001498 (iternextfunc)imap_next, /* tp_iternext */
1499 0, /* tp_methods */
1500 0, /* tp_members */
1501 0, /* tp_getset */
1502 0, /* tp_base */
1503 0, /* tp_dict */
1504 0, /* tp_descr_get */
1505 0, /* tp_descr_set */
1506 0, /* tp_dictoffset */
1507 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001508 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001509 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001510 PyObject_GC_Del, /* tp_free */
1511};
1512
1513
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001514/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001515
1516typedef struct {
1517 PyObject_HEAD
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001518 long tuplesize;
1519 long iternum; /* which iterator is active */
1520 PyObject *ittuple; /* tuple of iterators */
1521} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001522
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001523static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524
1525static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001526chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001527{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001528 chainobject *lz;
1529 int tuplesize = PySequence_Length(args);
1530 int i;
1531 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001532
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001533 /* obtain iterators */
1534 assert(PyTuple_Check(args));
1535 ittuple = PyTuple_New(tuplesize);
1536 if(ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001538 for (i=0; i < tuplesize; ++i) {
1539 PyObject *item = PyTuple_GET_ITEM(args, i);
1540 PyObject *it = PyObject_GetIter(item);
1541 if (it == NULL) {
1542 if (PyErr_ExceptionMatches(PyExc_TypeError))
1543 PyErr_Format(PyExc_TypeError,
1544 "chain argument #%d must support iteration",
1545 i+1);
1546 Py_DECREF(ittuple);
1547 return NULL;
1548 }
1549 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001550 }
1551
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001552 /* create chainobject structure */
1553 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001554 if (lz == NULL) {
1555 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001556 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001557 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001558
1559 lz->ittuple = ittuple;
1560 lz->iternum = 0;
1561 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562
1563 return (PyObject *)lz;
1564}
1565
1566static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001567chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001568{
1569 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001570 Py_XDECREF(lz->ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001571 lz->ob_type->tp_free(lz);
1572}
1573
Raymond Hettinger2012f172003-02-07 05:32:58 +00001574static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001575chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001576{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001577 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001578 return 0;
1579}
1580
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001581static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001582chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001584 PyObject *it;
1585 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001586
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001587 while (lz->iternum < lz->tuplesize) {
1588 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1589 item = PyIter_Next(it);
1590 if (item != NULL)
1591 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001592 if (PyErr_Occurred()) {
1593 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1594 PyErr_Clear();
1595 else
1596 return NULL;
1597 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001598 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001599 }
1600 return NULL;
1601}
1602
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001603PyDoc_STRVAR(chain_doc,
1604"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001605\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001606Return a chain object whose .next() method returns elements from the\n\
1607first iterable until it is exhausted, then elements from the next\n\
1608iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001610static PyTypeObject chain_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001611 PyObject_HEAD_INIT(NULL)
1612 0, /* ob_size */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001613 "itertools.chain", /* tp_name */
1614 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001615 0, /* tp_itemsize */
1616 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001617 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618 0, /* tp_print */
1619 0, /* tp_getattr */
1620 0, /* tp_setattr */
1621 0, /* tp_compare */
1622 0, /* tp_repr */
1623 0, /* tp_as_number */
1624 0, /* tp_as_sequence */
1625 0, /* tp_as_mapping */
1626 0, /* tp_hash */
1627 0, /* tp_call */
1628 0, /* tp_str */
1629 PyObject_GenericGetAttr, /* tp_getattro */
1630 0, /* tp_setattro */
1631 0, /* tp_as_buffer */
1632 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1633 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001634 chain_doc, /* tp_doc */
1635 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001636 0, /* tp_clear */
1637 0, /* tp_richcompare */
1638 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001639 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001640 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001641 0, /* tp_methods */
1642 0, /* tp_members */
1643 0, /* tp_getset */
1644 0, /* tp_base */
1645 0, /* tp_dict */
1646 0, /* tp_descr_get */
1647 0, /* tp_descr_set */
1648 0, /* tp_dictoffset */
1649 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001650 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001651 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001652 PyObject_GC_Del, /* tp_free */
1653};
1654
1655
1656/* ifilter object ************************************************************/
1657
1658typedef struct {
1659 PyObject_HEAD
1660 PyObject *func;
1661 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662} ifilterobject;
1663
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001664static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665
1666static PyObject *
1667ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1668{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001669 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670 PyObject *it;
1671 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001672
Raymond Hettinger60eca932003-02-09 06:40:58 +00001673 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674 return NULL;
1675
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676 /* Get iterator. */
1677 it = PyObject_GetIter(seq);
1678 if (it == NULL)
1679 return NULL;
1680
1681 /* create ifilterobject structure */
1682 lz = (ifilterobject *)type->tp_alloc(type, 0);
1683 if (lz == NULL) {
1684 Py_DECREF(it);
1685 return NULL;
1686 }
1687 Py_INCREF(func);
1688 lz->func = func;
1689 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001690
1691 return (PyObject *)lz;
1692}
1693
1694static void
1695ifilter_dealloc(ifilterobject *lz)
1696{
1697 PyObject_GC_UnTrack(lz);
1698 Py_XDECREF(lz->func);
1699 Py_XDECREF(lz->it);
1700 lz->ob_type->tp_free(lz);
1701}
1702
1703static int
1704ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1705{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001706 Py_VISIT(lz->it);
1707 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001708 return 0;
1709}
1710
1711static PyObject *
1712ifilter_next(ifilterobject *lz)
1713{
1714 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001715 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001717 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001718
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001719 assert(PyIter_Check(it));
1720 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001721 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001722 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001723 if (item == NULL)
1724 return NULL;
1725
1726 if (lz->func == Py_None) {
1727 ok = PyObject_IsTrue(item);
1728 } else {
1729 PyObject *good;
1730 good = PyObject_CallFunctionObjArgs(lz->func,
1731 item, NULL);
1732 if (good == NULL) {
1733 Py_DECREF(item);
1734 return NULL;
1735 }
1736 ok = PyObject_IsTrue(good);
1737 Py_DECREF(good);
1738 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001739 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001740 return item;
1741 Py_DECREF(item);
1742 }
1743}
1744
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001745PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001746"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001747\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001748Return those items of sequence for which function(item) is true.\n\
1749If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001750
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001751static PyTypeObject ifilter_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001752 PyObject_HEAD_INIT(NULL)
1753 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001754 "itertools.ifilter", /* tp_name */
1755 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001756 0, /* tp_itemsize */
1757 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001758 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759 0, /* tp_print */
1760 0, /* tp_getattr */
1761 0, /* tp_setattr */
1762 0, /* tp_compare */
1763 0, /* tp_repr */
1764 0, /* tp_as_number */
1765 0, /* tp_as_sequence */
1766 0, /* tp_as_mapping */
1767 0, /* tp_hash */
1768 0, /* tp_call */
1769 0, /* tp_str */
1770 PyObject_GenericGetAttr, /* tp_getattro */
1771 0, /* tp_setattro */
1772 0, /* tp_as_buffer */
1773 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1774 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001775 ifilter_doc, /* tp_doc */
1776 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001777 0, /* tp_clear */
1778 0, /* tp_richcompare */
1779 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001780 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001781 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001782 0, /* tp_methods */
1783 0, /* tp_members */
1784 0, /* tp_getset */
1785 0, /* tp_base */
1786 0, /* tp_dict */
1787 0, /* tp_descr_get */
1788 0, /* tp_descr_set */
1789 0, /* tp_dictoffset */
1790 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001791 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001792 ifilter_new, /* tp_new */
1793 PyObject_GC_Del, /* tp_free */
1794};
1795
1796
1797/* ifilterfalse object ************************************************************/
1798
1799typedef struct {
1800 PyObject_HEAD
1801 PyObject *func;
1802 PyObject *it;
1803} ifilterfalseobject;
1804
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001805static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001806
1807static PyObject *
1808ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1809{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001810 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001811 PyObject *it;
1812 ifilterfalseobject *lz;
1813
1814 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1815 return NULL;
1816
1817 /* Get iterator. */
1818 it = PyObject_GetIter(seq);
1819 if (it == NULL)
1820 return NULL;
1821
1822 /* create ifilterfalseobject structure */
1823 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1824 if (lz == NULL) {
1825 Py_DECREF(it);
1826 return NULL;
1827 }
1828 Py_INCREF(func);
1829 lz->func = func;
1830 lz->it = it;
1831
1832 return (PyObject *)lz;
1833}
1834
1835static void
1836ifilterfalse_dealloc(ifilterfalseobject *lz)
1837{
1838 PyObject_GC_UnTrack(lz);
1839 Py_XDECREF(lz->func);
1840 Py_XDECREF(lz->it);
1841 lz->ob_type->tp_free(lz);
1842}
1843
1844static int
1845ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1846{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001847 Py_VISIT(lz->it);
1848 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001849 return 0;
1850}
1851
1852static PyObject *
1853ifilterfalse_next(ifilterfalseobject *lz)
1854{
1855 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001856 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001857 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001858 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001859
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001860 assert(PyIter_Check(it));
1861 iternext = *it->ob_type->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001862 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001863 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001864 if (item == NULL)
1865 return NULL;
1866
1867 if (lz->func == Py_None) {
1868 ok = PyObject_IsTrue(item);
1869 } else {
1870 PyObject *good;
1871 good = PyObject_CallFunctionObjArgs(lz->func,
1872 item, NULL);
1873 if (good == NULL) {
1874 Py_DECREF(item);
1875 return NULL;
1876 }
1877 ok = PyObject_IsTrue(good);
1878 Py_DECREF(good);
1879 }
1880 if (!ok)
1881 return item;
1882 Py_DECREF(item);
1883 }
1884}
1885
Raymond Hettinger60eca932003-02-09 06:40:58 +00001886PyDoc_STRVAR(ifilterfalse_doc,
1887"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1888\n\
1889Return those items of sequence for which function(item) is false.\n\
1890If function is None, return the items that are false.");
1891
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001892static PyTypeObject ifilterfalse_type = {
Raymond Hettinger60eca932003-02-09 06:40:58 +00001893 PyObject_HEAD_INIT(NULL)
1894 0, /* ob_size */
1895 "itertools.ifilterfalse", /* tp_name */
1896 sizeof(ifilterfalseobject), /* tp_basicsize */
1897 0, /* tp_itemsize */
1898 /* methods */
1899 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1900 0, /* tp_print */
1901 0, /* tp_getattr */
1902 0, /* tp_setattr */
1903 0, /* tp_compare */
1904 0, /* tp_repr */
1905 0, /* tp_as_number */
1906 0, /* tp_as_sequence */
1907 0, /* tp_as_mapping */
1908 0, /* tp_hash */
1909 0, /* tp_call */
1910 0, /* tp_str */
1911 PyObject_GenericGetAttr, /* tp_getattro */
1912 0, /* tp_setattro */
1913 0, /* tp_as_buffer */
1914 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1915 Py_TPFLAGS_BASETYPE, /* tp_flags */
1916 ifilterfalse_doc, /* tp_doc */
1917 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
1918 0, /* tp_clear */
1919 0, /* tp_richcompare */
1920 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001921 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001922 (iternextfunc)ifilterfalse_next, /* tp_iternext */
1923 0, /* tp_methods */
1924 0, /* tp_members */
1925 0, /* tp_getset */
1926 0, /* tp_base */
1927 0, /* tp_dict */
1928 0, /* tp_descr_get */
1929 0, /* tp_descr_set */
1930 0, /* tp_dictoffset */
1931 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001932 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001933 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001934 PyObject_GC_Del, /* tp_free */
1935};
1936
1937
1938/* count object ************************************************************/
1939
1940typedef struct {
1941 PyObject_HEAD
1942 long cnt;
1943} countobject;
1944
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001945static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001946
1947static PyObject *
1948count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1949{
1950 countobject *lz;
1951 long cnt = 0;
1952
1953 if (!PyArg_ParseTuple(args, "|l:count", &cnt))
1954 return NULL;
1955
1956 /* create countobject structure */
1957 lz = (countobject *)PyObject_New(countobject, &count_type);
1958 if (lz == NULL)
1959 return NULL;
1960 lz->cnt = cnt;
1961
1962 return (PyObject *)lz;
1963}
1964
1965static PyObject *
1966count_next(countobject *lz)
1967{
1968 return PyInt_FromLong(lz->cnt++);
1969}
1970
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001971static PyObject *
1972count_repr(countobject *lz)
1973{
Brett Cannon00468392004-04-13 02:43:53 +00001974 return PyString_FromFormat("count(%ld)", lz->cnt);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001975}
1976
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001977PyDoc_STRVAR(count_doc,
1978"count([firstval]) --> count object\n\
1979\n\
1980Return a count object whose .next() method returns consecutive\n\
1981integers starting from zero or, if specified, from firstval.");
1982
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001983static PyTypeObject count_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001984 PyObject_HEAD_INIT(NULL)
1985 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001986 "itertools.count", /* tp_name */
1987 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001988 0, /* tp_itemsize */
1989 /* methods */
1990 (destructor)PyObject_Del, /* tp_dealloc */
1991 0, /* tp_print */
1992 0, /* tp_getattr */
1993 0, /* tp_setattr */
1994 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00001995 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001996 0, /* tp_as_number */
1997 0, /* tp_as_sequence */
1998 0, /* tp_as_mapping */
1999 0, /* tp_hash */
2000 0, /* tp_call */
2001 0, /* tp_str */
2002 PyObject_GenericGetAttr, /* tp_getattro */
2003 0, /* tp_setattro */
2004 0, /* tp_as_buffer */
2005 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002006 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002007 0, /* tp_traverse */
2008 0, /* tp_clear */
2009 0, /* tp_richcompare */
2010 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002011 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002012 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002013 0, /* tp_methods */
2014 0, /* tp_members */
2015 0, /* tp_getset */
2016 0, /* tp_base */
2017 0, /* tp_dict */
2018 0, /* tp_descr_get */
2019 0, /* tp_descr_set */
2020 0, /* tp_dictoffset */
2021 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002022 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002023 count_new, /* tp_new */
2024};
2025
2026
2027/* izip object ************************************************************/
2028
2029#include "Python.h"
2030
2031typedef struct {
2032 PyObject_HEAD
2033 long tuplesize;
2034 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002035 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002036} izipobject;
2037
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002038static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002039
2040static PyObject *
2041izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2042{
2043 izipobject *lz;
2044 int i;
2045 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002046 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002047 int tuplesize = PySequence_Length(args);
2048
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002049 /* args must be a tuple */
2050 assert(PyTuple_Check(args));
2051
2052 /* obtain iterators */
2053 ittuple = PyTuple_New(tuplesize);
2054 if(ittuple == NULL)
2055 return NULL;
2056 for (i=0; i < tuplesize; ++i) {
2057 PyObject *item = PyTuple_GET_ITEM(args, i);
2058 PyObject *it = PyObject_GetIter(item);
2059 if (it == NULL) {
2060 if (PyErr_ExceptionMatches(PyExc_TypeError))
2061 PyErr_Format(PyExc_TypeError,
2062 "izip argument #%d must support iteration",
2063 i+1);
2064 Py_DECREF(ittuple);
2065 return NULL;
2066 }
2067 PyTuple_SET_ITEM(ittuple, i, it);
2068 }
2069
Raymond Hettinger2012f172003-02-07 05:32:58 +00002070 /* create a result holder */
2071 result = PyTuple_New(tuplesize);
2072 if (result == NULL) {
2073 Py_DECREF(ittuple);
2074 return NULL;
2075 }
2076 for (i=0 ; i < tuplesize ; i++) {
2077 Py_INCREF(Py_None);
2078 PyTuple_SET_ITEM(result, i, Py_None);
2079 }
2080
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002081 /* create izipobject structure */
2082 lz = (izipobject *)type->tp_alloc(type, 0);
2083 if (lz == NULL) {
2084 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002085 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002086 return NULL;
2087 }
2088 lz->ittuple = ittuple;
2089 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002090 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002091
2092 return (PyObject *)lz;
2093}
2094
2095static void
2096izip_dealloc(izipobject *lz)
2097{
2098 PyObject_GC_UnTrack(lz);
2099 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002100 Py_XDECREF(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002101 lz->ob_type->tp_free(lz);
2102}
2103
2104static int
2105izip_traverse(izipobject *lz, visitproc visit, void *arg)
2106{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002107 Py_VISIT(lz->ittuple);
2108 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002109 return 0;
2110}
2111
2112static PyObject *
2113izip_next(izipobject *lz)
2114{
2115 int i;
2116 long tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002117 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002118 PyObject *it;
2119 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002120 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002121
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002122 if (tuplesize == 0)
2123 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002124 if (result->ob_refcnt == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002125 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002126 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002127 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002128 assert(PyIter_Check(it));
2129 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002130 if (item == NULL) {
2131 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002132 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002133 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002134 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002135 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002136 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002137 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002138 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002139 result = PyTuple_New(tuplesize);
2140 if (result == NULL)
2141 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002142 for (i=0 ; i < tuplesize ; i++) {
2143 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002144 assert(PyIter_Check(it));
2145 item = (*it->ob_type->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002146 if (item == NULL) {
2147 Py_DECREF(result);
2148 return NULL;
2149 }
2150 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002151 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002152 }
2153 return result;
2154}
2155
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002156PyDoc_STRVAR(izip_doc,
2157"izip(iter1 [,iter2 [...]]) --> izip object\n\
2158\n\
2159Return a izip object whose .next() method returns a tuple where\n\
2160the i-th element comes from the i-th iterable argument. The .next()\n\
2161method continues until the shortest iterable in the argument sequence\n\
2162is exhausted and then it raises StopIteration. Works like the zip()\n\
2163function but consumes less memory by returning an iterator instead of\n\
2164a list.");
2165
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002166static PyTypeObject izip_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002167 PyObject_HEAD_INIT(NULL)
2168 0, /* ob_size */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002169 "itertools.izip", /* tp_name */
2170 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002171 0, /* tp_itemsize */
2172 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002173 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002174 0, /* tp_print */
2175 0, /* tp_getattr */
2176 0, /* tp_setattr */
2177 0, /* tp_compare */
2178 0, /* tp_repr */
2179 0, /* tp_as_number */
2180 0, /* tp_as_sequence */
2181 0, /* tp_as_mapping */
2182 0, /* tp_hash */
2183 0, /* tp_call */
2184 0, /* tp_str */
2185 PyObject_GenericGetAttr, /* tp_getattro */
2186 0, /* tp_setattro */
2187 0, /* tp_as_buffer */
2188 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2189 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002190 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002191 (traverseproc)izip_traverse, /* tp_traverse */
2192 0, /* tp_clear */
2193 0, /* tp_richcompare */
2194 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002195 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002196 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002197 0, /* tp_methods */
2198 0, /* tp_members */
2199 0, /* tp_getset */
2200 0, /* tp_base */
2201 0, /* tp_dict */
2202 0, /* tp_descr_get */
2203 0, /* tp_descr_set */
2204 0, /* tp_dictoffset */
2205 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002206 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002207 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002208 PyObject_GC_Del, /* tp_free */
2209};
2210
2211
2212/* repeat object ************************************************************/
2213
2214typedef struct {
2215 PyObject_HEAD
2216 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002217 long cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002218} repeatobject;
2219
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002220static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002221
2222static PyObject *
2223repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2224{
2225 repeatobject *ro;
2226 PyObject *element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002227 long cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002228
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002229 if (!PyArg_ParseTuple(args, "O|l:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002230 return NULL;
2231
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002232 if (PyTuple_Size(args) == 2 && cnt < 0)
2233 cnt = 0;
2234
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002235 ro = (repeatobject *)type->tp_alloc(type, 0);
2236 if (ro == NULL)
2237 return NULL;
2238 Py_INCREF(element);
2239 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002240 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002241 return (PyObject *)ro;
2242}
2243
2244static void
2245repeat_dealloc(repeatobject *ro)
2246{
2247 PyObject_GC_UnTrack(ro);
2248 Py_XDECREF(ro->element);
2249 ro->ob_type->tp_free(ro);
2250}
2251
2252static int
2253repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2254{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002255 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002256 return 0;
2257}
2258
2259static PyObject *
2260repeat_next(repeatobject *ro)
2261{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002262 if (ro->cnt == 0)
2263 return NULL;
2264 if (ro->cnt > 0)
2265 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002266 Py_INCREF(ro->element);
2267 return ro->element;
2268}
2269
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002270static PyObject *
2271repeat_repr(repeatobject *ro)
2272{
2273 PyObject *result, *objrepr;
2274
2275 objrepr = PyObject_Repr(ro->element);
2276 if (objrepr == NULL)
2277 return NULL;
2278
2279 if (ro->cnt == -1)
2280 result = PyString_FromFormat("repeat(%s)",
2281 PyString_AS_STRING(objrepr));
2282 else
Brett Cannon00468392004-04-13 02:43:53 +00002283 result = PyString_FromFormat("repeat(%s, %ld)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002284 PyString_AS_STRING(objrepr), ro->cnt);
2285 Py_DECREF(objrepr);
2286 return result;
2287}
2288
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002289static int
2290repeat_len(repeatobject *ro)
2291{
2292 if (ro->cnt == -1)
2293 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
2294 return (int)(ro->cnt);
2295}
2296
2297static PySequenceMethods repeat_as_sequence = {
2298 (inquiry)repeat_len, /* sq_length */
2299 0, /* sq_concat */
2300};
2301
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002302PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002303"repeat(element [,times]) -> create an iterator which returns the element\n\
2304for the specified number of times. If not specified, returns the element\n\
2305endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002306
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002307static PyTypeObject repeat_type = {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002308 PyObject_HEAD_INIT(NULL)
2309 0, /* ob_size */
2310 "itertools.repeat", /* tp_name */
2311 sizeof(repeatobject), /* tp_basicsize */
2312 0, /* tp_itemsize */
2313 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002314 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002315 0, /* tp_print */
2316 0, /* tp_getattr */
2317 0, /* tp_setattr */
2318 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002319 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002320 0, /* tp_as_number */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002321 &repeat_as_sequence, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002322 0, /* tp_as_mapping */
2323 0, /* tp_hash */
2324 0, /* tp_call */
2325 0, /* tp_str */
2326 PyObject_GenericGetAttr, /* tp_getattro */
2327 0, /* tp_setattro */
2328 0, /* tp_as_buffer */
2329 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2330 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002331 repeat_doc, /* tp_doc */
2332 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002333 0, /* tp_clear */
2334 0, /* tp_richcompare */
2335 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002336 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002337 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002338 0, /* tp_methods */
2339 0, /* tp_members */
2340 0, /* tp_getset */
2341 0, /* tp_base */
2342 0, /* tp_dict */
2343 0, /* tp_descr_get */
2344 0, /* tp_descr_set */
2345 0, /* tp_dictoffset */
2346 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002347 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002348 repeat_new, /* tp_new */
2349 PyObject_GC_Del, /* tp_free */
2350};
2351
2352
2353/* module level code ********************************************************/
2354
2355PyDoc_STRVAR(module_doc,
2356"Functional tools for creating and using iterators.\n\
2357\n\
2358Infinite iterators:\n\
2359count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002360cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002361repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002362\n\
2363Iterators terminating on the shortest input sequence:\n\
2364izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002365ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2366ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002367islice(seq, [start,] stop [, step]) --> elements from\n\
2368 seq[start:stop:step]\n\
2369imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2370starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002371tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002372chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002373takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2374dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002375groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002376");
2377
2378
Raymond Hettingerad983e72003-11-12 14:32:26 +00002379static PyMethodDef module_methods[] = {
2380 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2381 {NULL, NULL} /* sentinel */
2382};
2383
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002384PyMODINIT_FUNC
2385inititertools(void)
2386{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002387 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002388 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002389 char *name;
2390 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002391 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002392 &dropwhile_type,
2393 &takewhile_type,
2394 &islice_type,
2395 &starmap_type,
2396 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002397 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002398 &ifilter_type,
2399 &ifilterfalse_type,
2400 &count_type,
2401 &izip_type,
2402 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002403 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002404 NULL
2405 };
2406
Skip Montanarof3938fd2004-02-10 20:27:40 +00002407 teedataobject_type.ob_type = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002408 m = Py_InitModule3("itertools", module_methods, module_doc);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002409
Raymond Hettinger60eca932003-02-09 06:40:58 +00002410 for (i=0 ; typelist[i] != NULL ; i++) {
2411 if (PyType_Ready(typelist[i]) < 0)
2412 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002413 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002414 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002415 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002416 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002417 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002418
2419 if (PyType_Ready(&teedataobject_type) < 0)
2420 return;
2421 if (PyType_Ready(&tee_type) < 0)
2422 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002423 if (PyType_Ready(&_grouper_type) < 0)
2424 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002425}