blob: fb54b1413e8b0784e28d7d678d1ab49797c7be08 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
2#include "Python.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00003#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004
5/* Itertools module written and maintained
6 by Raymond D. Hettinger <python@rcn.com>
7 Copyright (c) 2003 Python Software Foundation.
8 All rights reserved.
9*/
10
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000011
12/* groupby object ***********************************************************/
13
14typedef struct {
15 PyObject_HEAD
16 PyObject *it;
17 PyObject *keyfunc;
18 PyObject *tgtkey;
19 PyObject *currkey;
20 PyObject *currvalue;
21} groupbyobject;
22
23static PyTypeObject groupby_type;
24static PyObject *_grouper_create(groupbyobject *, PyObject *);
25
26static PyObject *
27groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
28{
Martin v. Löwis02cbf4a2006-02-27 17:20:04 +000029 static char *kwargs[] = {"iterable", "key", NULL};
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000030 groupbyobject *gbo;
31 PyObject *it, *keyfunc = Py_None;
32
33 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,
34 &it, &keyfunc))
35 return NULL;
36
37 gbo = (groupbyobject *)type->tp_alloc(type, 0);
38 if (gbo == NULL)
39 return NULL;
40 gbo->tgtkey = NULL;
41 gbo->currkey = NULL;
42 gbo->currvalue = NULL;
43 gbo->keyfunc = keyfunc;
44 Py_INCREF(keyfunc);
45 gbo->it = PyObject_GetIter(it);
46 if (gbo->it == NULL) {
47 Py_DECREF(gbo);
48 return NULL;
49 }
50 return (PyObject *)gbo;
51}
52
53static void
54groupby_dealloc(groupbyobject *gbo)
55{
56 PyObject_GC_UnTrack(gbo);
57 Py_XDECREF(gbo->it);
58 Py_XDECREF(gbo->keyfunc);
59 Py_XDECREF(gbo->tgtkey);
60 Py_XDECREF(gbo->currkey);
61 Py_XDECREF(gbo->currvalue);
Martin v. Löwis68192102007-07-21 06:55:02 +000062 Py_Type(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063}
64
65static int
66groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
67{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +000068 Py_VISIT(gbo->it);
69 Py_VISIT(gbo->keyfunc);
70 Py_VISIT(gbo->tgtkey);
71 Py_VISIT(gbo->currkey);
72 Py_VISIT(gbo->currvalue);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000073 return 0;
74}
75
76static PyObject *
77groupby_next(groupbyobject *gbo)
78{
Raymond Hettinger4cda01e2004-09-28 04:45:28 +000079 PyObject *newvalue, *newkey, *r, *grouper, *tmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000080
81 /* skip to next iteration group */
82 for (;;) {
83 if (gbo->currkey == NULL)
84 /* pass */;
85 else if (gbo->tgtkey == NULL)
86 break;
87 else {
88 int rcmp;
89
90 rcmp = PyObject_RichCompareBool(gbo->tgtkey,
91 gbo->currkey, Py_EQ);
92 if (rcmp == -1)
93 return NULL;
94 else if (rcmp == 0)
95 break;
96 }
97
98 newvalue = PyIter_Next(gbo->it);
99 if (newvalue == NULL)
100 return NULL;
101
102 if (gbo->keyfunc == Py_None) {
103 newkey = newvalue;
104 Py_INCREF(newvalue);
105 } else {
106 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
107 newvalue, NULL);
108 if (newkey == NULL) {
109 Py_DECREF(newvalue);
110 return NULL;
111 }
112 }
113
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000114 tmp = gbo->currkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000115 gbo->currkey = newkey;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000116 Py_XDECREF(tmp);
117
118 tmp = gbo->currvalue;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000119 gbo->currvalue = newvalue;
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000120 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000121 }
122
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000123 Py_INCREF(gbo->currkey);
Raymond Hettinger4cda01e2004-09-28 04:45:28 +0000124 tmp = gbo->tgtkey;
125 gbo->tgtkey = gbo->currkey;
126 Py_XDECREF(tmp);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000127
128 grouper = _grouper_create(gbo, gbo->tgtkey);
129 if (grouper == NULL)
130 return NULL;
131
132 r = PyTuple_Pack(2, gbo->currkey, grouper);
133 Py_DECREF(grouper);
134 return r;
135}
136
137PyDoc_STRVAR(groupby_doc,
138"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\
139(key, sub-iterator) grouped by each value of key(value).\n");
140
141static PyTypeObject groupby_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000142 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000143 "itertools.groupby", /* tp_name */
144 sizeof(groupbyobject), /* tp_basicsize */
145 0, /* tp_itemsize */
146 /* methods */
147 (destructor)groupby_dealloc, /* tp_dealloc */
148 0, /* tp_print */
149 0, /* tp_getattr */
150 0, /* tp_setattr */
151 0, /* tp_compare */
152 0, /* tp_repr */
153 0, /* tp_as_number */
154 0, /* tp_as_sequence */
155 0, /* tp_as_mapping */
156 0, /* tp_hash */
157 0, /* tp_call */
158 0, /* tp_str */
159 PyObject_GenericGetAttr, /* tp_getattro */
160 0, /* tp_setattro */
161 0, /* tp_as_buffer */
162 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
163 Py_TPFLAGS_BASETYPE, /* tp_flags */
164 groupby_doc, /* tp_doc */
165 (traverseproc)groupby_traverse, /* tp_traverse */
166 0, /* tp_clear */
167 0, /* tp_richcompare */
168 0, /* tp_weaklistoffset */
169 PyObject_SelfIter, /* tp_iter */
170 (iternextfunc)groupby_next, /* tp_iternext */
171 0, /* tp_methods */
172 0, /* tp_members */
173 0, /* tp_getset */
174 0, /* tp_base */
175 0, /* tp_dict */
176 0, /* tp_descr_get */
177 0, /* tp_descr_set */
178 0, /* tp_dictoffset */
179 0, /* tp_init */
180 0, /* tp_alloc */
181 groupby_new, /* tp_new */
182 PyObject_GC_Del, /* tp_free */
183};
184
185
186/* _grouper object (internal) ************************************************/
187
188typedef struct {
189 PyObject_HEAD
190 PyObject *parent;
191 PyObject *tgtkey;
192} _grouperobject;
193
194static PyTypeObject _grouper_type;
195
196static PyObject *
197_grouper_create(groupbyobject *parent, PyObject *tgtkey)
198{
199 _grouperobject *igo;
200
201 igo = PyObject_New(_grouperobject, &_grouper_type);
202 if (igo == NULL)
203 return NULL;
204 igo->parent = (PyObject *)parent;
205 Py_INCREF(parent);
206 igo->tgtkey = tgtkey;
207 Py_INCREF(tgtkey);
208
209 return (PyObject *)igo;
210}
211
212static void
213_grouper_dealloc(_grouperobject *igo)
214{
215 Py_DECREF(igo->parent);
216 Py_DECREF(igo->tgtkey);
217 PyObject_Del(igo);
218}
219
220static PyObject *
221_grouper_next(_grouperobject *igo)
222{
223 groupbyobject *gbo = (groupbyobject *)igo->parent;
224 PyObject *newvalue, *newkey, *r;
225 int rcmp;
226
227 if (gbo->currvalue == NULL) {
228 newvalue = PyIter_Next(gbo->it);
229 if (newvalue == NULL)
230 return NULL;
231
232 if (gbo->keyfunc == Py_None) {
233 newkey = newvalue;
234 Py_INCREF(newvalue);
235 } else {
236 newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,
237 newvalue, NULL);
238 if (newkey == NULL) {
239 Py_DECREF(newvalue);
240 return NULL;
241 }
242 }
243
244 assert(gbo->currkey == NULL);
245 gbo->currkey = newkey;
246 gbo->currvalue = newvalue;
247 }
248
249 assert(gbo->currkey != NULL);
250 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
251 if (rcmp <= 0)
252 /* got any error or current group is end */
253 return NULL;
254
255 r = gbo->currvalue;
256 gbo->currvalue = NULL;
Raymond Hettinger75ccea32004-09-01 07:02:44 +0000257 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000258
259 return r;
260}
261
262static PyTypeObject _grouper_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000263 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000264 "itertools._grouper", /* tp_name */
265 sizeof(_grouperobject), /* tp_basicsize */
266 0, /* tp_itemsize */
267 /* methods */
268 (destructor)_grouper_dealloc, /* tp_dealloc */
269 0, /* tp_print */
270 0, /* tp_getattr */
271 0, /* tp_setattr */
272 0, /* tp_compare */
273 0, /* tp_repr */
274 0, /* tp_as_number */
275 0, /* tp_as_sequence */
276 0, /* tp_as_mapping */
277 0, /* tp_hash */
278 0, /* tp_call */
279 0, /* tp_str */
280 PyObject_GenericGetAttr, /* tp_getattro */
281 0, /* tp_setattro */
282 0, /* tp_as_buffer */
283 Py_TPFLAGS_DEFAULT, /* tp_flags */
284 0, /* tp_doc */
285 0, /* tp_traverse */
286 0, /* tp_clear */
287 0, /* tp_richcompare */
288 0, /* tp_weaklistoffset */
289 PyObject_SelfIter, /* tp_iter */
290 (iternextfunc)_grouper_next, /* tp_iternext */
291 0, /* tp_methods */
292 0, /* tp_members */
293 0, /* tp_getset */
294 0, /* tp_base */
295 0, /* tp_dict */
296 0, /* tp_descr_get */
297 0, /* tp_descr_set */
298 0, /* tp_dictoffset */
299 0, /* tp_init */
300 0, /* tp_alloc */
301 0, /* tp_new */
302 PyObject_Del, /* tp_free */
303};
304
305
306
Raymond Hettingerad983e72003-11-12 14:32:26 +0000307/* tee object and with supporting function and objects ***************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000308
Raymond Hettingerad983e72003-11-12 14:32:26 +0000309/* The teedataobject pre-allocates space for LINKCELLS number of objects.
310 To help the object fit neatly inside cache lines (space for 16 to 32
311 pointers), the value should be a multiple of 16 minus space for
312 the other structure members including PyHEAD overhead. The larger the
313 value, the less memory overhead per object and the less time spent
314 allocating/deallocating new links. The smaller the number, the less
315 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000316*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000317#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000318
319typedef struct {
320 PyObject_HEAD
321 PyObject *it;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000322 int numread;
323 PyObject *nextlink;
324 PyObject *(values[LINKCELLS]);
325} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000326
327typedef struct {
328 PyObject_HEAD
Raymond Hettingerad983e72003-11-12 14:32:26 +0000329 teedataobject *dataobj;
330 int index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000331 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000332} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000333
Raymond Hettingerad983e72003-11-12 14:32:26 +0000334static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335
336static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000337teedataobject_new(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000338{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000339 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000340
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000341 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000342 if (tdo == NULL)
Raymond Hettinger45143692003-10-25 06:37:47 +0000343 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000344
345 tdo->numread = 0;
346 tdo->nextlink = NULL;
347 Py_INCREF(it);
348 tdo->it = it;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000349 PyObject_GC_Track(tdo);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000350 return (PyObject *)tdo;
351}
352
353static PyObject *
354teedataobject_jumplink(teedataobject *tdo)
355{
356 if (tdo->nextlink == NULL)
357 tdo->nextlink = teedataobject_new(tdo->it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000358 Py_XINCREF(tdo->nextlink);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000359 return tdo->nextlink;
360}
361
362static PyObject *
363teedataobject_getitem(teedataobject *tdo, int i)
364{
365 PyObject *value;
366
367 assert(i < LINKCELLS);
368 if (i < tdo->numread)
369 value = tdo->values[i];
370 else {
371 /* this is the lead iterator, so fetch more data */
372 assert(i == tdo->numread);
373 value = PyIter_Next(tdo->it);
374 if (value == NULL)
375 return NULL;
376 tdo->numread++;
377 tdo->values[i] = value;
Raymond Hettinger45143692003-10-25 06:37:47 +0000378 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000379 Py_INCREF(value);
380 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000381}
382
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000383static int
384teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
385{
386 int i;
387 Py_VISIT(tdo->it);
388 for (i = 0; i < tdo->numread; i++)
389 Py_VISIT(tdo->values[i]);
390 Py_VISIT(tdo->nextlink);
391 return 0;
392}
393
394static int
395teedataobject_clear(teedataobject *tdo)
396{
397 int i;
398 Py_CLEAR(tdo->it);
399 for (i=0 ; i<tdo->numread ; i++)
400 Py_CLEAR(tdo->values[i]);
401 Py_CLEAR(tdo->nextlink);
402 return 0;
403}
404
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000405static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000406teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000407{
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000408 PyObject_GC_UnTrack(tdo);
409 teedataobject_clear(tdo);
410 PyObject_GC_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 = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000416 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000417 "itertools.tee_dataobject", /* tp_name */
418 sizeof(teedataobject), /* tp_basicsize */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000419 0, /* tp_itemsize */
420 /* methods */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000421 (destructor)teedataobject_dealloc, /* tp_dealloc */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000422 0, /* tp_print */
423 0, /* tp_getattr */
424 0, /* tp_setattr */
425 0, /* tp_compare */
426 0, /* tp_repr */
427 0, /* tp_as_number */
428 0, /* tp_as_sequence */
429 0, /* tp_as_mapping */
430 0, /* tp_hash */
431 0, /* tp_call */
432 0, /* tp_str */
433 PyObject_GenericGetAttr, /* tp_getattro */
434 0, /* tp_setattro */
435 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000437 teedataobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000438 (traverseproc)teedataobject_traverse, /* tp_traverse */
439 (inquiry)teedataobject_clear, /* tp_clear */
440 0, /* tp_richcompare */
441 0, /* tp_weaklistoffset */
442 0, /* tp_iter */
443 0, /* tp_iternext */
444 0, /* tp_methods */
445 0, /* tp_members */
446 0, /* tp_getset */
447 0, /* tp_base */
448 0, /* tp_dict */
449 0, /* tp_descr_get */
450 0, /* tp_descr_set */
451 0, /* tp_dictoffset */
452 0, /* tp_init */
453 0, /* tp_alloc */
454 0, /* tp_new */
455 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000456};
457
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000458
459static PyTypeObject tee_type;
460
461static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000462tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463{
Raymond Hettingerad983e72003-11-12 14:32:26 +0000464 PyObject *value, *link;
465
466 if (to->index >= LINKCELLS) {
467 link = teedataobject_jumplink(to->dataobj);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000468 Py_DECREF(to->dataobj);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000469 to->dataobj = (teedataobject *)link;
470 to->index = 0;
471 }
472 value = teedataobject_getitem(to->dataobj, to->index);
473 if (value == NULL)
474 return NULL;
475 to->index++;
476 return value;
477}
478
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000479static int
480tee_traverse(teeobject *to, visitproc visit, void *arg)
481{
482 Py_VISIT((PyObject *)to->dataobj);
483 return 0;
484}
485
Raymond Hettingerad983e72003-11-12 14:32:26 +0000486static PyObject *
487tee_copy(teeobject *to)
488{
489 teeobject *newto;
490
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000491 newto = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000492 if (newto == NULL)
493 return NULL;
494 Py_INCREF(to->dataobj);
495 newto->dataobj = to->dataobj;
496 newto->index = to->index;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000497 newto->weakreflist = NULL;
Thomas Woutersb3deb942006-04-15 22:33:13 +0000498 PyObject_GC_Track(newto);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000499 return (PyObject *)newto;
500}
501
502PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
503
504static PyObject *
505tee_fromiterable(PyObject *iterable)
506{
507 teeobject *to;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000508 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000509
510 it = PyObject_GetIter(iterable);
511 if (it == NULL)
512 return NULL;
513 if (PyObject_TypeCheck(it, &tee_type)) {
514 to = (teeobject *)tee_copy((teeobject *)it);
515 goto done;
516 }
517
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000518 to = PyObject_GC_New(teeobject, &tee_type);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000519 if (to == NULL)
520 goto done;
521 to->dataobj = (teedataobject *)teedataobject_new(it);
Neal Norwitz9029b5f2006-07-23 07:59:00 +0000522 if (!to->dataobj) {
523 PyObject_GC_Del(to);
524 to = NULL;
525 goto done;
526 }
527
Raymond Hettingerad983e72003-11-12 14:32:26 +0000528 to->index = 0;
Raymond Hettingera9f60922004-10-17 16:40:14 +0000529 to->weakreflist = NULL;
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000530 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000531done:
532 Py_XDECREF(it);
533 return (PyObject *)to;
534}
535
536static PyObject *
537tee_new(PyTypeObject *type, PyObject *args, PyObject *kw)
538{
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000539 PyObject *iterable;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000540
541 if (!PyArg_UnpackTuple(args, "tee", 1, 1, &iterable))
542 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000543 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000544}
545
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000546static int
547tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000548{
Raymond Hettingera9f60922004-10-17 16:40:14 +0000549 if (to->weakreflist != NULL)
550 PyObject_ClearWeakRefs((PyObject *) to);
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000551 Py_CLEAR(to->dataobj);
552 return 0;
553}
554
555static void
556tee_dealloc(teeobject *to)
557{
558 PyObject_GC_UnTrack(to);
559 tee_clear(to);
560 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000561}
562
Raymond Hettingerad983e72003-11-12 14:32:26 +0000563PyDoc_STRVAR(teeobject_doc,
564"Iterator wrapped to make it copyable");
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000565
Raymond Hettingerad983e72003-11-12 14:32:26 +0000566static PyMethodDef tee_methods[] = {
567 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
568 {NULL, NULL} /* sentinel */
569};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000570
571static PyTypeObject tee_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000572 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000573 "itertools.tee", /* tp_name */
574 sizeof(teeobject), /* tp_basicsize */
575 0, /* tp_itemsize */
576 /* methods */
577 (destructor)tee_dealloc, /* tp_dealloc */
578 0, /* tp_print */
579 0, /* tp_getattr */
580 0, /* tp_setattr */
581 0, /* tp_compare */
582 0, /* tp_repr */
583 0, /* tp_as_number */
584 0, /* tp_as_sequence */
585 0, /* tp_as_mapping */
586 0, /* tp_hash */
587 0, /* tp_call */
588 0, /* tp_str */
Raymond Hettingerf0c5aec2003-10-26 14:25:56 +0000589 0, /* tp_getattro */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000590 0, /* tp_setattro */
591 0, /* tp_as_buffer */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000592 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000593 teeobject_doc, /* tp_doc */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000594 (traverseproc)tee_traverse, /* tp_traverse */
595 (inquiry)tee_clear, /* tp_clear */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000596 0, /* tp_richcompare */
Raymond Hettingera9f60922004-10-17 16:40:14 +0000597 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000598 PyObject_SelfIter, /* tp_iter */
599 (iternextfunc)tee_next, /* tp_iternext */
600 tee_methods, /* tp_methods */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000601 0, /* tp_members */
602 0, /* tp_getset */
603 0, /* tp_base */
604 0, /* tp_dict */
605 0, /* tp_descr_get */
606 0, /* tp_descr_set */
607 0, /* tp_dictoffset */
608 0, /* tp_init */
609 0, /* tp_alloc */
610 tee_new, /* tp_new */
Thomas Wouters19bf33b2006-03-27 21:02:13 +0000611 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000612};
613
Raymond Hettingerad983e72003-11-12 14:32:26 +0000614static PyObject *
615tee(PyObject *self, PyObject *args)
616{
Neal Norwitz69e88972006-09-02 02:58:13 +0000617 Py_ssize_t i, n=2;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000618 PyObject *it, *iterable, *copyable, *result;
619
Neal Norwitz69e88972006-09-02 02:58:13 +0000620 if (!PyArg_ParseTuple(args, "O|n", &iterable, &n))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000621 return NULL;
Neal Norwitz69e88972006-09-02 02:58:13 +0000622 if (n < 0) {
623 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
624 return NULL;
625 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000626 result = PyTuple_New(n);
627 if (result == NULL)
628 return NULL;
629 if (n == 0)
630 return result;
631 it = PyObject_GetIter(iterable);
632 if (it == NULL) {
633 Py_DECREF(result);
634 return NULL;
635 }
636 if (!PyObject_HasAttrString(it, "__copy__")) {
637 copyable = tee_fromiterable(it);
638 Py_DECREF(it);
639 if (copyable == NULL) {
640 Py_DECREF(result);
641 return NULL;
642 }
643 } else
644 copyable = it;
645 PyTuple_SET_ITEM(result, 0, copyable);
646 for (i=1 ; i<n ; i++) {
647 copyable = PyObject_CallMethod(copyable, "__copy__", NULL);
648 if (copyable == NULL) {
649 Py_DECREF(result);
650 return NULL;
651 }
652 PyTuple_SET_ITEM(result, i, copyable);
653 }
654 return result;
655}
656
657PyDoc_STRVAR(tee_doc,
658"tee(iterable, n=2) --> tuple of n independent iterators.");
659
660
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000661/* cycle object **********************************************************/
662
663typedef struct {
664 PyObject_HEAD
665 PyObject *it;
666 PyObject *saved;
667 int firstpass;
668} cycleobject;
669
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000670static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000671
672static PyObject *
673cycle_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
674{
675 PyObject *it;
676 PyObject *iterable;
677 PyObject *saved;
678 cycleobject *lz;
679
Georg Brandlb84c1372007-01-21 10:28:43 +0000680 if (type == &cycle_type && !_PyArg_NoKeywords("cycle()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000681 return NULL;
682
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000683 if (!PyArg_UnpackTuple(args, "cycle", 1, 1, &iterable))
684 return NULL;
685
686 /* Get iterator. */
687 it = PyObject_GetIter(iterable);
688 if (it == NULL)
689 return NULL;
690
691 saved = PyList_New(0);
692 if (saved == NULL) {
693 Py_DECREF(it);
694 return NULL;
695 }
696
697 /* create cycleobject structure */
698 lz = (cycleobject *)type->tp_alloc(type, 0);
699 if (lz == NULL) {
700 Py_DECREF(it);
701 Py_DECREF(saved);
702 return NULL;
703 }
704 lz->it = it;
705 lz->saved = saved;
706 lz->firstpass = 0;
707
708 return (PyObject *)lz;
709}
710
711static void
712cycle_dealloc(cycleobject *lz)
713{
714 PyObject_GC_UnTrack(lz);
715 Py_XDECREF(lz->saved);
716 Py_XDECREF(lz->it);
Martin v. Löwis68192102007-07-21 06:55:02 +0000717 Py_Type(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000718}
719
720static int
721cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
722{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000723 Py_VISIT(lz->it);
724 Py_VISIT(lz->saved);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000725 return 0;
726}
727
728static PyObject *
729cycle_next(cycleobject *lz)
730{
731 PyObject *item;
732 PyObject *it;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000733 PyObject *tmp;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000734
735 while (1) {
736 item = PyIter_Next(lz->it);
737 if (item != NULL) {
738 if (!lz->firstpass)
739 PyList_Append(lz->saved, item);
740 return item;
741 }
Raymond Hettinger9d7c8702004-05-08 19:49:42 +0000742 if (PyErr_Occurred()) {
743 if (PyErr_ExceptionMatches(PyExc_StopIteration))
744 PyErr_Clear();
745 else
746 return NULL;
747 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000748 if (PyList_Size(lz->saved) == 0)
749 return NULL;
750 it = PyObject_GetIter(lz->saved);
751 if (it == NULL)
752 return NULL;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000753 tmp = lz->it;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000754 lz->it = it;
755 lz->firstpass = 1;
Raymond Hettinger880430e2004-10-02 10:56:43 +0000756 Py_DECREF(tmp);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000757 }
758}
759
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000760PyDoc_STRVAR(cycle_doc,
761"cycle(iterable) --> cycle object\n\
762\n\
763Return elements from the iterable until it is exhausted.\n\
764Then repeat the sequence indefinitely.");
765
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000766static PyTypeObject cycle_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000767 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000768 "itertools.cycle", /* tp_name */
769 sizeof(cycleobject), /* tp_basicsize */
770 0, /* tp_itemsize */
771 /* methods */
772 (destructor)cycle_dealloc, /* tp_dealloc */
773 0, /* tp_print */
774 0, /* tp_getattr */
775 0, /* tp_setattr */
776 0, /* tp_compare */
777 0, /* tp_repr */
778 0, /* tp_as_number */
779 0, /* tp_as_sequence */
780 0, /* tp_as_mapping */
781 0, /* tp_hash */
782 0, /* tp_call */
783 0, /* tp_str */
784 PyObject_GenericGetAttr, /* tp_getattro */
785 0, /* tp_setattro */
786 0, /* tp_as_buffer */
787 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
788 Py_TPFLAGS_BASETYPE, /* tp_flags */
789 cycle_doc, /* tp_doc */
790 (traverseproc)cycle_traverse, /* tp_traverse */
791 0, /* tp_clear */
792 0, /* tp_richcompare */
793 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000794 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000795 (iternextfunc)cycle_next, /* tp_iternext */
796 0, /* tp_methods */
797 0, /* tp_members */
798 0, /* tp_getset */
799 0, /* tp_base */
800 0, /* tp_dict */
801 0, /* tp_descr_get */
802 0, /* tp_descr_set */
803 0, /* tp_dictoffset */
804 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000805 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000806 cycle_new, /* tp_new */
807 PyObject_GC_Del, /* tp_free */
808};
809
810
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000811/* dropwhile object **********************************************************/
812
813typedef struct {
814 PyObject_HEAD
815 PyObject *func;
816 PyObject *it;
817 long start;
818} dropwhileobject;
819
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000820static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000821
822static PyObject *
823dropwhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
824{
825 PyObject *func, *seq;
826 PyObject *it;
827 dropwhileobject *lz;
828
Georg Brandlb84c1372007-01-21 10:28:43 +0000829 if (type == &dropwhile_type && !_PyArg_NoKeywords("dropwhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000830 return NULL;
831
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000832 if (!PyArg_UnpackTuple(args, "dropwhile", 2, 2, &func, &seq))
833 return NULL;
834
835 /* Get iterator. */
836 it = PyObject_GetIter(seq);
837 if (it == NULL)
838 return NULL;
839
840 /* create dropwhileobject structure */
841 lz = (dropwhileobject *)type->tp_alloc(type, 0);
842 if (lz == NULL) {
843 Py_DECREF(it);
844 return NULL;
845 }
846 Py_INCREF(func);
847 lz->func = func;
848 lz->it = it;
849 lz->start = 0;
850
851 return (PyObject *)lz;
852}
853
854static void
855dropwhile_dealloc(dropwhileobject *lz)
856{
857 PyObject_GC_UnTrack(lz);
858 Py_XDECREF(lz->func);
859 Py_XDECREF(lz->it);
Martin v. Löwis68192102007-07-21 06:55:02 +0000860 Py_Type(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000861}
862
863static int
864dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
865{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +0000866 Py_VISIT(lz->it);
867 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000868 return 0;
869}
870
871static PyObject *
872dropwhile_next(dropwhileobject *lz)
873{
874 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +0000875 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000876 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000877 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000878
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000879 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +0000880 iternext = *Py_Type(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000881 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +0000882 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000883 if (item == NULL)
884 return NULL;
885 if (lz->start == 1)
886 return item;
887
888 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
889 if (good == NULL) {
890 Py_DECREF(item);
891 return NULL;
892 }
893 ok = PyObject_IsTrue(good);
894 Py_DECREF(good);
895 if (!ok) {
896 lz->start = 1;
897 return item;
898 }
899 Py_DECREF(item);
900 }
901}
902
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000903PyDoc_STRVAR(dropwhile_doc,
904"dropwhile(predicate, iterable) --> dropwhile object\n\
905\n\
906Drop items from the iterable while predicate(item) is true.\n\
907Afterwards, return every element until the iterable is exhausted.");
908
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000909static PyTypeObject dropwhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +0000910 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000911 "itertools.dropwhile", /* tp_name */
912 sizeof(dropwhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000913 0, /* tp_itemsize */
914 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000915 (destructor)dropwhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000916 0, /* tp_print */
917 0, /* tp_getattr */
918 0, /* tp_setattr */
919 0, /* tp_compare */
920 0, /* tp_repr */
921 0, /* tp_as_number */
922 0, /* tp_as_sequence */
923 0, /* tp_as_mapping */
924 0, /* tp_hash */
925 0, /* tp_call */
926 0, /* tp_str */
927 PyObject_GenericGetAttr, /* tp_getattro */
928 0, /* tp_setattro */
929 0, /* tp_as_buffer */
930 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
931 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000932 dropwhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000933 (traverseproc)dropwhile_traverse, /* tp_traverse */
934 0, /* tp_clear */
935 0, /* tp_richcompare */
936 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +0000937 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000938 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000939 0, /* tp_methods */
940 0, /* tp_members */
941 0, /* tp_getset */
942 0, /* tp_base */
943 0, /* tp_dict */
944 0, /* tp_descr_get */
945 0, /* tp_descr_set */
946 0, /* tp_dictoffset */
947 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +0000948 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +0000949 dropwhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000950 PyObject_GC_Del, /* tp_free */
951};
952
953
954/* takewhile object **********************************************************/
955
956typedef struct {
957 PyObject_HEAD
958 PyObject *func;
959 PyObject *it;
960 long stop;
961} takewhileobject;
962
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000963static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000964
965static PyObject *
966takewhile_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
967{
968 PyObject *func, *seq;
969 PyObject *it;
970 takewhileobject *lz;
971
Georg Brandlb84c1372007-01-21 10:28:43 +0000972 if (type == &takewhile_type && !_PyArg_NoKeywords("takewhile()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +0000973 return NULL;
974
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000975 if (!PyArg_UnpackTuple(args, "takewhile", 2, 2, &func, &seq))
976 return NULL;
977
978 /* Get iterator. */
979 it = PyObject_GetIter(seq);
980 if (it == NULL)
981 return NULL;
982
983 /* create takewhileobject structure */
984 lz = (takewhileobject *)type->tp_alloc(type, 0);
985 if (lz == NULL) {
986 Py_DECREF(it);
987 return NULL;
988 }
989 Py_INCREF(func);
990 lz->func = func;
991 lz->it = it;
992 lz->stop = 0;
993
994 return (PyObject *)lz;
995}
996
997static void
998takewhile_dealloc(takewhileobject *lz)
999{
1000 PyObject_GC_UnTrack(lz);
1001 Py_XDECREF(lz->func);
1002 Py_XDECREF(lz->it);
Martin v. Löwis68192102007-07-21 06:55:02 +00001003 Py_Type(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001004}
1005
1006static int
1007takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1008{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001009 Py_VISIT(lz->it);
1010 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001011 return 0;
1012}
1013
1014static PyObject *
1015takewhile_next(takewhileobject *lz)
1016{
1017 PyObject *item, *good;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001018 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001019 long ok;
1020
1021 if (lz->stop == 1)
1022 return NULL;
1023
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001024 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00001025 item = (*Py_Type(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001026 if (item == NULL)
1027 return NULL;
1028
1029 good = PyObject_CallFunctionObjArgs(lz->func, item, NULL);
1030 if (good == NULL) {
1031 Py_DECREF(item);
1032 return NULL;
1033 }
1034 ok = PyObject_IsTrue(good);
1035 Py_DECREF(good);
1036 if (ok)
1037 return item;
1038 Py_DECREF(item);
1039 lz->stop = 1;
1040 return NULL;
1041}
1042
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001043PyDoc_STRVAR(takewhile_doc,
1044"takewhile(predicate, iterable) --> takewhile object\n\
1045\n\
1046Return successive entries from an iterable as long as the \n\
1047predicate evaluates to true for each entry.");
1048
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001049static PyTypeObject takewhile_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001050 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001051 "itertools.takewhile", /* tp_name */
1052 sizeof(takewhileobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001053 0, /* tp_itemsize */
1054 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001055 (destructor)takewhile_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001056 0, /* tp_print */
1057 0, /* tp_getattr */
1058 0, /* tp_setattr */
1059 0, /* tp_compare */
1060 0, /* tp_repr */
1061 0, /* tp_as_number */
1062 0, /* tp_as_sequence */
1063 0, /* tp_as_mapping */
1064 0, /* tp_hash */
1065 0, /* tp_call */
1066 0, /* tp_str */
1067 PyObject_GenericGetAttr, /* tp_getattro */
1068 0, /* tp_setattro */
1069 0, /* tp_as_buffer */
1070 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1071 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001072 takewhile_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001073 (traverseproc)takewhile_traverse, /* tp_traverse */
1074 0, /* tp_clear */
1075 0, /* tp_richcompare */
1076 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001077 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001078 (iternextfunc)takewhile_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001079 0, /* tp_methods */
1080 0, /* tp_members */
1081 0, /* tp_getset */
1082 0, /* tp_base */
1083 0, /* tp_dict */
1084 0, /* tp_descr_get */
1085 0, /* tp_descr_set */
1086 0, /* tp_dictoffset */
1087 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001088 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001089 takewhile_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001090 PyObject_GC_Del, /* tp_free */
1091};
1092
1093
1094/* islice object ************************************************************/
1095
1096typedef struct {
1097 PyObject_HEAD
1098 PyObject *it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001099 Py_ssize_t next;
1100 Py_ssize_t stop;
1101 Py_ssize_t step;
1102 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001103} isliceobject;
1104
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001105static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001106
1107static PyObject *
1108islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1109{
1110 PyObject *seq;
Jack Diederich6c433a92006-05-26 11:15:17 +00001111 Py_ssize_t start=0, stop=-1, step=1;
Raymond Hettingerb2594052004-12-05 09:25:51 +00001112 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001113 Py_ssize_t numargs;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001114 isliceobject *lz;
1115
Georg Brandlb84c1372007-01-21 10:28:43 +00001116 if (type == &islice_type && !_PyArg_NoKeywords("islice()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001117 return NULL;
1118
Raymond Hettingerb2594052004-12-05 09:25:51 +00001119 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001120 return NULL;
1121
Raymond Hettingerb2594052004-12-05 09:25:51 +00001122 numargs = PyTuple_Size(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001123 if (numargs == 2) {
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001124 if (a1 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001125 stop = PyInt_AsSsize_t(a1);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001126 if (stop == -1) {
1127 if (PyErr_Occurred())
1128 PyErr_Clear();
1129 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001130 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001131 return NULL;
1132 }
1133 }
Raymond Hettinger341deb72003-05-02 19:44:20 +00001134 } else {
Raymond Hettingerb2594052004-12-05 09:25:51 +00001135 if (a1 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001136 start = PyInt_AsSsize_t(a1);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001137 if (start == -1 && PyErr_Occurred())
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001138 PyErr_Clear();
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001139 if (a2 != Py_None) {
Jack Diederich6c433a92006-05-26 11:15:17 +00001140 stop = PyInt_AsSsize_t(a2);
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001141 if (stop == -1) {
1142 if (PyErr_Occurred())
1143 PyErr_Clear();
1144 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001145 "Stop argument for islice() must be a non-negative integer or None.");
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001146 return NULL;
1147 }
1148 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001149 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001150 if (start<0 || stop<-1) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001151 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001152 "Indices for islice() must be non-negative integers or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001153 return NULL;
1154 }
1155
Raymond Hettingerb2594052004-12-05 09:25:51 +00001156 if (a3 != NULL) {
1157 if (a3 != Py_None)
Jack Diederich6c433a92006-05-26 11:15:17 +00001158 step = PyInt_AsSsize_t(a3);
Raymond Hettingerb2594052004-12-05 09:25:51 +00001159 if (step == -1 && PyErr_Occurred())
1160 PyErr_Clear();
1161 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001162 if (step<1) {
1163 PyErr_SetString(PyExc_ValueError,
Raymond Hettingerb2594052004-12-05 09:25:51 +00001164 "Step for islice() must be a positive integer or None.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165 return NULL;
1166 }
1167
1168 /* Get iterator. */
1169 it = PyObject_GetIter(seq);
1170 if (it == NULL)
1171 return NULL;
1172
1173 /* create isliceobject structure */
1174 lz = (isliceobject *)type->tp_alloc(type, 0);
1175 if (lz == NULL) {
1176 Py_DECREF(it);
1177 return NULL;
1178 }
1179 lz->it = it;
1180 lz->next = start;
1181 lz->stop = stop;
1182 lz->step = step;
1183 lz->cnt = 0L;
1184
1185 return (PyObject *)lz;
1186}
1187
1188static void
1189islice_dealloc(isliceobject *lz)
1190{
1191 PyObject_GC_UnTrack(lz);
1192 Py_XDECREF(lz->it);
Martin v. Löwis68192102007-07-21 06:55:02 +00001193 Py_Type(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001194}
1195
1196static int
1197islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1198{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001199 Py_VISIT(lz->it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001200 return 0;
1201}
1202
1203static PyObject *
1204islice_next(isliceobject *lz)
1205{
1206 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001207 PyObject *it = lz->it;
Jack Diederich6c433a92006-05-26 11:15:17 +00001208 Py_ssize_t oldnext;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001209 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001210
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001211 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00001212 iternext = *Py_Type(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213 while (lz->cnt < lz->next) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001214 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001215 if (item == NULL)
1216 return NULL;
1217 Py_DECREF(item);
1218 lz->cnt++;
1219 }
Raymond Hettinger14ef54c2003-05-02 19:04:37 +00001220 if (lz->stop != -1 && lz->cnt >= lz->stop)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221 return NULL;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001222 assert(PyIter_Check(it));
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001223 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001224 if (item == NULL)
1225 return NULL;
1226 lz->cnt++;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001227 oldnext = lz->next;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228 lz->next += lz->step;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001229 if (lz->next < oldnext) /* Check for overflow */
1230 lz->next = lz->stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001231 return item;
1232}
1233
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001234PyDoc_STRVAR(islice_doc,
1235"islice(iterable, [start,] stop [, step]) --> islice object\n\
1236\n\
1237Return an iterator whose next() method returns selected values from an\n\
1238iterable. If start is specified, will skip all preceding elements;\n\
1239otherwise, start defaults to zero. Step defaults to one. If\n\
1240specified as another value, step determines how many values are \n\
1241skipped between successive calls. Works like a slice() on a list\n\
1242but returns an iterator.");
1243
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001244static PyTypeObject islice_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001245 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001246 "itertools.islice", /* tp_name */
1247 sizeof(isliceobject), /* tp_basicsize */
1248 0, /* tp_itemsize */
1249 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001250 (destructor)islice_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001251 0, /* tp_print */
1252 0, /* tp_getattr */
1253 0, /* tp_setattr */
1254 0, /* tp_compare */
1255 0, /* tp_repr */
1256 0, /* tp_as_number */
1257 0, /* tp_as_sequence */
1258 0, /* tp_as_mapping */
1259 0, /* tp_hash */
1260 0, /* tp_call */
1261 0, /* tp_str */
1262 PyObject_GenericGetAttr, /* tp_getattro */
1263 0, /* tp_setattro */
1264 0, /* tp_as_buffer */
1265 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1266 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001267 islice_doc, /* tp_doc */
1268 (traverseproc)islice_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001269 0, /* tp_clear */
1270 0, /* tp_richcompare */
1271 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001272 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001273 (iternextfunc)islice_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001274 0, /* tp_methods */
1275 0, /* tp_members */
1276 0, /* tp_getset */
1277 0, /* tp_base */
1278 0, /* tp_dict */
1279 0, /* tp_descr_get */
1280 0, /* tp_descr_set */
1281 0, /* tp_dictoffset */
1282 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001283 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001284 islice_new, /* tp_new */
1285 PyObject_GC_Del, /* tp_free */
1286};
1287
1288
1289/* starmap object ************************************************************/
1290
1291typedef struct {
1292 PyObject_HEAD
1293 PyObject *func;
1294 PyObject *it;
1295} starmapobject;
1296
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001297static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001298
1299static PyObject *
1300starmap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1301{
1302 PyObject *func, *seq;
1303 PyObject *it;
1304 starmapobject *lz;
1305
Georg Brandlb84c1372007-01-21 10:28:43 +00001306 if (type == &starmap_type && !_PyArg_NoKeywords("starmap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001307 return NULL;
1308
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001309 if (!PyArg_UnpackTuple(args, "starmap", 2, 2, &func, &seq))
1310 return NULL;
1311
1312 /* Get iterator. */
1313 it = PyObject_GetIter(seq);
1314 if (it == NULL)
1315 return NULL;
1316
1317 /* create starmapobject structure */
1318 lz = (starmapobject *)type->tp_alloc(type, 0);
1319 if (lz == NULL) {
1320 Py_DECREF(it);
1321 return NULL;
1322 }
1323 Py_INCREF(func);
1324 lz->func = func;
1325 lz->it = it;
1326
1327 return (PyObject *)lz;
1328}
1329
1330static void
1331starmap_dealloc(starmapobject *lz)
1332{
1333 PyObject_GC_UnTrack(lz);
1334 Py_XDECREF(lz->func);
1335 Py_XDECREF(lz->it);
Martin v. Löwis68192102007-07-21 06:55:02 +00001336 Py_Type(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001337}
1338
1339static int
1340starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1341{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001342 Py_VISIT(lz->it);
1343 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001344 return 0;
1345}
1346
1347static PyObject *
1348starmap_next(starmapobject *lz)
1349{
1350 PyObject *args;
1351 PyObject *result;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001352 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001353
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001354 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00001355 args = (*Py_Type(it)->tp_iternext)(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001356 if (args == NULL)
1357 return NULL;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001358 if (!PyTuple_CheckExact(args)) {
1359 Py_DECREF(args);
1360 PyErr_SetString(PyExc_TypeError,
1361 "iterator must return a tuple");
1362 return NULL;
1363 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001364 result = PyObject_Call(lz->func, args, NULL);
1365 Py_DECREF(args);
1366 return result;
1367}
1368
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001369PyDoc_STRVAR(starmap_doc,
1370"starmap(function, sequence) --> starmap object\n\
1371\n\
1372Return an iterator whose values are returned from the function evaluated\n\
1373with a argument tuple taken from the given sequence.");
1374
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001375static PyTypeObject starmap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001376 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001377 "itertools.starmap", /* tp_name */
1378 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001379 0, /* tp_itemsize */
1380 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001381 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001382 0, /* tp_print */
1383 0, /* tp_getattr */
1384 0, /* tp_setattr */
1385 0, /* tp_compare */
1386 0, /* tp_repr */
1387 0, /* tp_as_number */
1388 0, /* tp_as_sequence */
1389 0, /* tp_as_mapping */
1390 0, /* tp_hash */
1391 0, /* tp_call */
1392 0, /* tp_str */
1393 PyObject_GenericGetAttr, /* tp_getattro */
1394 0, /* tp_setattro */
1395 0, /* tp_as_buffer */
1396 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1397 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001398 starmap_doc, /* tp_doc */
1399 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001400 0, /* tp_clear */
1401 0, /* tp_richcompare */
1402 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001403 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001404 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001405 0, /* tp_methods */
1406 0, /* tp_members */
1407 0, /* tp_getset */
1408 0, /* tp_base */
1409 0, /* tp_dict */
1410 0, /* tp_descr_get */
1411 0, /* tp_descr_set */
1412 0, /* tp_dictoffset */
1413 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001414 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001415 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001416 PyObject_GC_Del, /* tp_free */
1417};
1418
1419
1420/* imap object ************************************************************/
1421
1422typedef struct {
1423 PyObject_HEAD
1424 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001425 PyObject *func;
1426} imapobject;
1427
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001428static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001429
1430static PyObject *
1431imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1432{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001433 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001434 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001435 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001436
Georg Brandlb84c1372007-01-21 10:28:43 +00001437 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001438 return NULL;
1439
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440 numargs = PyTuple_Size(args);
1441 if (numargs < 2) {
1442 PyErr_SetString(PyExc_TypeError,
1443 "imap() must have at least two arguments.");
1444 return NULL;
1445 }
1446
1447 iters = PyTuple_New(numargs-1);
1448 if (iters == NULL)
1449 return NULL;
1450
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001451 for (i=1 ; i<numargs ; i++) {
1452 /* Get iterator. */
1453 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1454 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001455 Py_DECREF(iters);
1456 return NULL;
1457 }
1458 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001459 }
1460
1461 /* create imapobject structure */
1462 lz = (imapobject *)type->tp_alloc(type, 0);
1463 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001464 Py_DECREF(iters);
1465 return NULL;
1466 }
1467 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468 func = PyTuple_GET_ITEM(args, 0);
1469 Py_INCREF(func);
1470 lz->func = func;
1471
1472 return (PyObject *)lz;
1473}
1474
1475static void
1476imap_dealloc(imapobject *lz)
1477{
1478 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479 Py_XDECREF(lz->iters);
1480 Py_XDECREF(lz->func);
Martin v. Löwis68192102007-07-21 06:55:02 +00001481 Py_Type(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482}
1483
1484static int
1485imap_traverse(imapobject *lz, visitproc visit, void *arg)
1486{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001487 Py_VISIT(lz->iters);
1488 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001489 return 0;
1490}
1491
Raymond Hettinger2012f172003-02-07 05:32:58 +00001492/*
1493imap() is an iterator version of __builtins__.map() except that it does
1494not have the None fill-in feature. That was intentionally left out for
1495the following reasons:
1496
1497 1) Itertools are designed to be easily combined and chained together.
1498 Having all tools stop with the shortest input is a unifying principle
1499 that makes it easier to combine finite iterators (supplying data) with
1500 infinite iterators like count() and repeat() (for supplying sequential
1501 or constant arguments to a function).
1502
1503 2) In typical use cases for combining itertools, having one finite data
1504 supplier run out before another is likely to be an error condition which
1505 should not pass silently by automatically supplying None.
1506
1507 3) The use cases for automatic None fill-in are rare -- not many functions
1508 do something useful when a parameter suddenly switches type and becomes
1509 None.
1510
1511 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001512 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001513
1514 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1515*/
1516
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001517static PyObject *
1518imap_next(imapobject *lz)
1519{
1520 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001521 PyObject *argtuple;
1522 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001523 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524
1525 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001526 argtuple = PyTuple_New(numargs);
1527 if (argtuple == NULL)
1528 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001529
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001530 for (i=0 ; i<numargs ; i++) {
1531 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1532 if (val == NULL) {
1533 Py_DECREF(argtuple);
1534 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001536 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001537 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001538 if (lz->func == Py_None)
1539 return argtuple;
1540 result = PyObject_Call(lz->func, argtuple, NULL);
1541 Py_DECREF(argtuple);
1542 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001543}
1544
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001545PyDoc_STRVAR(imap_doc,
1546"imap(func, *iterables) --> imap object\n\
1547\n\
1548Make an iterator that computes the function using arguments from\n\
1549each of the iterables. Like map() except that it returns\n\
1550an iterator instead of a list and that it stops when the shortest\n\
1551iterable is exhausted instead of filling in None for shorter\n\
1552iterables.");
1553
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001554static PyTypeObject imap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001555 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001556 "itertools.imap", /* tp_name */
1557 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001558 0, /* tp_itemsize */
1559 /* methods */
1560 (destructor)imap_dealloc, /* tp_dealloc */
1561 0, /* tp_print */
1562 0, /* tp_getattr */
1563 0, /* tp_setattr */
1564 0, /* tp_compare */
1565 0, /* tp_repr */
1566 0, /* tp_as_number */
1567 0, /* tp_as_sequence */
1568 0, /* tp_as_mapping */
1569 0, /* tp_hash */
1570 0, /* tp_call */
1571 0, /* tp_str */
1572 PyObject_GenericGetAttr, /* tp_getattro */
1573 0, /* tp_setattro */
1574 0, /* tp_as_buffer */
1575 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1576 Py_TPFLAGS_BASETYPE, /* tp_flags */
1577 imap_doc, /* tp_doc */
1578 (traverseproc)imap_traverse, /* tp_traverse */
1579 0, /* tp_clear */
1580 0, /* tp_richcompare */
1581 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001582 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001583 (iternextfunc)imap_next, /* tp_iternext */
1584 0, /* tp_methods */
1585 0, /* tp_members */
1586 0, /* tp_getset */
1587 0, /* tp_base */
1588 0, /* tp_dict */
1589 0, /* tp_descr_get */
1590 0, /* tp_descr_set */
1591 0, /* tp_dictoffset */
1592 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001593 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001594 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001595 PyObject_GC_Del, /* tp_free */
1596};
1597
1598
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001599/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001600
1601typedef struct {
1602 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00001603 Py_ssize_t tuplesize;
1604 Py_ssize_t iternum; /* which iterator is active */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001605 PyObject *ittuple; /* tuple of iterators */
1606} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001607
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001608static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609
1610static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001611chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001612{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001613 chainobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001614 Py_ssize_t tuplesize = PySequence_Length(args);
Jack Diederich6c433a92006-05-26 11:15:17 +00001615 Py_ssize_t i;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001616 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001617
Georg Brandlb84c1372007-01-21 10:28:43 +00001618 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001619 return NULL;
1620
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001621 /* obtain iterators */
1622 assert(PyTuple_Check(args));
1623 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001624 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001625 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001626 for (i=0; i < tuplesize; ++i) {
1627 PyObject *item = PyTuple_GET_ITEM(args, i);
1628 PyObject *it = PyObject_GetIter(item);
1629 if (it == NULL) {
1630 if (PyErr_ExceptionMatches(PyExc_TypeError))
1631 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001632 "chain argument #%zd must support iteration",
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001633 i+1);
1634 Py_DECREF(ittuple);
1635 return NULL;
1636 }
1637 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001638 }
1639
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001640 /* create chainobject structure */
1641 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001642 if (lz == NULL) {
1643 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001644 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001645 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001646
1647 lz->ittuple = ittuple;
1648 lz->iternum = 0;
1649 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001650
1651 return (PyObject *)lz;
1652}
1653
1654static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001655chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001656{
1657 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001658 Py_XDECREF(lz->ittuple);
Martin v. Löwis68192102007-07-21 06:55:02 +00001659 Py_Type(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001660}
1661
Raymond Hettinger2012f172003-02-07 05:32:58 +00001662static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001663chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001664{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001665 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001666 return 0;
1667}
1668
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001669static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001670chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001672 PyObject *it;
1673 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001675 while (lz->iternum < lz->tuplesize) {
1676 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1677 item = PyIter_Next(it);
1678 if (item != NULL)
1679 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001680 if (PyErr_Occurred()) {
1681 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1682 PyErr_Clear();
1683 else
1684 return NULL;
1685 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001686 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001687 }
1688 return NULL;
1689}
1690
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001691PyDoc_STRVAR(chain_doc,
1692"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001693\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001694Return a chain object whose .next() method returns elements from the\n\
1695first iterable until it is exhausted, then elements from the next\n\
1696iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001697
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001698static PyTypeObject chain_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001699 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001700 "itertools.chain", /* tp_name */
1701 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001702 0, /* tp_itemsize */
1703 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001704 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001705 0, /* tp_print */
1706 0, /* tp_getattr */
1707 0, /* tp_setattr */
1708 0, /* tp_compare */
1709 0, /* tp_repr */
1710 0, /* tp_as_number */
1711 0, /* tp_as_sequence */
1712 0, /* tp_as_mapping */
1713 0, /* tp_hash */
1714 0, /* tp_call */
1715 0, /* tp_str */
1716 PyObject_GenericGetAttr, /* tp_getattro */
1717 0, /* tp_setattro */
1718 0, /* tp_as_buffer */
1719 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1720 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001721 chain_doc, /* tp_doc */
1722 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001723 0, /* tp_clear */
1724 0, /* tp_richcompare */
1725 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001726 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001727 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001728 0, /* tp_methods */
1729 0, /* tp_members */
1730 0, /* tp_getset */
1731 0, /* tp_base */
1732 0, /* tp_dict */
1733 0, /* tp_descr_get */
1734 0, /* tp_descr_set */
1735 0, /* tp_dictoffset */
1736 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001737 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001738 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739 PyObject_GC_Del, /* tp_free */
1740};
1741
1742
1743/* ifilter object ************************************************************/
1744
1745typedef struct {
1746 PyObject_HEAD
1747 PyObject *func;
1748 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001749} ifilterobject;
1750
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001751static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001752
1753static PyObject *
1754ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1755{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001756 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757 PyObject *it;
1758 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759
Georg Brandlb84c1372007-01-21 10:28:43 +00001760 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001761 return NULL;
1762
Raymond Hettinger60eca932003-02-09 06:40:58 +00001763 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001764 return NULL;
1765
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001766 /* Get iterator. */
1767 it = PyObject_GetIter(seq);
1768 if (it == NULL)
1769 return NULL;
1770
1771 /* create ifilterobject structure */
1772 lz = (ifilterobject *)type->tp_alloc(type, 0);
1773 if (lz == NULL) {
1774 Py_DECREF(it);
1775 return NULL;
1776 }
1777 Py_INCREF(func);
1778 lz->func = func;
1779 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001780
1781 return (PyObject *)lz;
1782}
1783
1784static void
1785ifilter_dealloc(ifilterobject *lz)
1786{
1787 PyObject_GC_UnTrack(lz);
1788 Py_XDECREF(lz->func);
1789 Py_XDECREF(lz->it);
Martin v. Löwis68192102007-07-21 06:55:02 +00001790 Py_Type(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001791}
1792
1793static int
1794ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
1795{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001796 Py_VISIT(lz->it);
1797 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001798 return 0;
1799}
1800
1801static PyObject *
1802ifilter_next(ifilterobject *lz)
1803{
1804 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001805 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001806 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001807 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001808
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001809 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00001810 iternext = *Py_Type(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001811 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001812 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001813 if (item == NULL)
1814 return NULL;
1815
1816 if (lz->func == Py_None) {
1817 ok = PyObject_IsTrue(item);
1818 } else {
1819 PyObject *good;
1820 good = PyObject_CallFunctionObjArgs(lz->func,
1821 item, NULL);
1822 if (good == NULL) {
1823 Py_DECREF(item);
1824 return NULL;
1825 }
1826 ok = PyObject_IsTrue(good);
1827 Py_DECREF(good);
1828 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00001829 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001830 return item;
1831 Py_DECREF(item);
1832 }
1833}
1834
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001835PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00001836"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001837\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00001838Return those items of sequence for which function(item) is true.\n\
1839If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001840
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001841static PyTypeObject ifilter_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001842 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001843 "itertools.ifilter", /* tp_name */
1844 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001845 0, /* tp_itemsize */
1846 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001847 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001848 0, /* tp_print */
1849 0, /* tp_getattr */
1850 0, /* tp_setattr */
1851 0, /* tp_compare */
1852 0, /* tp_repr */
1853 0, /* tp_as_number */
1854 0, /* tp_as_sequence */
1855 0, /* tp_as_mapping */
1856 0, /* tp_hash */
1857 0, /* tp_call */
1858 0, /* tp_str */
1859 PyObject_GenericGetAttr, /* tp_getattro */
1860 0, /* tp_setattro */
1861 0, /* tp_as_buffer */
1862 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1863 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001864 ifilter_doc, /* tp_doc */
1865 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001866 0, /* tp_clear */
1867 0, /* tp_richcompare */
1868 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001869 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001870 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001871 0, /* tp_methods */
1872 0, /* tp_members */
1873 0, /* tp_getset */
1874 0, /* tp_base */
1875 0, /* tp_dict */
1876 0, /* tp_descr_get */
1877 0, /* tp_descr_set */
1878 0, /* tp_dictoffset */
1879 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001880 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001881 ifilter_new, /* tp_new */
1882 PyObject_GC_Del, /* tp_free */
1883};
1884
1885
1886/* ifilterfalse object ************************************************************/
1887
1888typedef struct {
1889 PyObject_HEAD
1890 PyObject *func;
1891 PyObject *it;
1892} ifilterfalseobject;
1893
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001894static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001895
1896static PyObject *
1897ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1898{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00001899 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001900 PyObject *it;
1901 ifilterfalseobject *lz;
1902
Georg Brandlb84c1372007-01-21 10:28:43 +00001903 if (type == &ifilterfalse_type &&
1904 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001905 return NULL;
1906
Raymond Hettinger60eca932003-02-09 06:40:58 +00001907 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
1908 return NULL;
1909
1910 /* Get iterator. */
1911 it = PyObject_GetIter(seq);
1912 if (it == NULL)
1913 return NULL;
1914
1915 /* create ifilterfalseobject structure */
1916 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
1917 if (lz == NULL) {
1918 Py_DECREF(it);
1919 return NULL;
1920 }
1921 Py_INCREF(func);
1922 lz->func = func;
1923 lz->it = it;
1924
1925 return (PyObject *)lz;
1926}
1927
1928static void
1929ifilterfalse_dealloc(ifilterfalseobject *lz)
1930{
1931 PyObject_GC_UnTrack(lz);
1932 Py_XDECREF(lz->func);
1933 Py_XDECREF(lz->it);
Martin v. Löwis68192102007-07-21 06:55:02 +00001934 Py_Type(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001935}
1936
1937static int
1938ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
1939{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001940 Py_VISIT(lz->it);
1941 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001942 return 0;
1943}
1944
1945static PyObject *
1946ifilterfalse_next(ifilterfalseobject *lz)
1947{
1948 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00001949 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001950 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001951 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001952
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001953 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00001954 iternext = *Py_Type(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00001955 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00001956 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00001957 if (item == NULL)
1958 return NULL;
1959
1960 if (lz->func == Py_None) {
1961 ok = PyObject_IsTrue(item);
1962 } else {
1963 PyObject *good;
1964 good = PyObject_CallFunctionObjArgs(lz->func,
1965 item, NULL);
1966 if (good == NULL) {
1967 Py_DECREF(item);
1968 return NULL;
1969 }
1970 ok = PyObject_IsTrue(good);
1971 Py_DECREF(good);
1972 }
1973 if (!ok)
1974 return item;
1975 Py_DECREF(item);
1976 }
1977}
1978
Raymond Hettinger60eca932003-02-09 06:40:58 +00001979PyDoc_STRVAR(ifilterfalse_doc,
1980"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
1981\n\
1982Return those items of sequence for which function(item) is false.\n\
1983If function is None, return the items that are false.");
1984
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001985static PyTypeObject ifilterfalse_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001986 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001987 "itertools.ifilterfalse", /* tp_name */
1988 sizeof(ifilterfalseobject), /* tp_basicsize */
1989 0, /* tp_itemsize */
1990 /* methods */
1991 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
1992 0, /* tp_print */
1993 0, /* tp_getattr */
1994 0, /* tp_setattr */
1995 0, /* tp_compare */
1996 0, /* tp_repr */
1997 0, /* tp_as_number */
1998 0, /* tp_as_sequence */
1999 0, /* tp_as_mapping */
2000 0, /* tp_hash */
2001 0, /* tp_call */
2002 0, /* tp_str */
2003 PyObject_GenericGetAttr, /* tp_getattro */
2004 0, /* tp_setattro */
2005 0, /* tp_as_buffer */
2006 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2007 Py_TPFLAGS_BASETYPE, /* tp_flags */
2008 ifilterfalse_doc, /* tp_doc */
2009 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2010 0, /* tp_clear */
2011 0, /* tp_richcompare */
2012 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002013 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002014 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2015 0, /* tp_methods */
2016 0, /* tp_members */
2017 0, /* tp_getset */
2018 0, /* tp_base */
2019 0, /* tp_dict */
2020 0, /* tp_descr_get */
2021 0, /* tp_descr_set */
2022 0, /* tp_dictoffset */
2023 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002024 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002025 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002026 PyObject_GC_Del, /* tp_free */
2027};
2028
2029
2030/* count object ************************************************************/
2031
2032typedef struct {
2033 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002034 Py_ssize_t cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002035 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002036} countobject;
2037
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002038static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002039
2040static PyObject *
2041count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2042{
2043 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002044 Py_ssize_t cnt = 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002045 PyObject *cnt_arg = NULL;
2046 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002047
Georg Brandlb84c1372007-01-21 10:28:43 +00002048 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002049 return NULL;
2050
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002051 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002052 return NULL;
2053
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002054 if (cnt_arg != NULL) {
2055 cnt = PyInt_AsSsize_t(cnt_arg);
2056 if (cnt == -1 && PyErr_Occurred()) {
2057 PyErr_Clear();
2058 if (!PyLong_Check(cnt_arg)) {
2059 PyErr_SetString(PyExc_TypeError, "an integer is required");
2060 return NULL;
2061 }
2062 long_cnt = cnt_arg;
2063 Py_INCREF(long_cnt);
2064 cnt = PY_SSIZE_T_MAX;
2065 }
2066 }
2067
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002068 /* create countobject structure */
2069 lz = (countobject *)PyObject_New(countobject, &count_type);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002070 if (lz == NULL) {
2071 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002072 return NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002073 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002074 lz->cnt = cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002075 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002076
2077 return (PyObject *)lz;
2078}
2079
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002080static void
2081count_dealloc(countobject *lz)
2082{
2083 Py_XDECREF(lz->long_cnt);
2084 PyObject_Del(lz);
2085}
2086
2087static PyObject *
2088count_nextlong(countobject *lz)
2089{
2090 static PyObject *one = NULL;
2091 PyObject *cnt;
2092 PyObject *stepped_up;
2093
2094 if (lz->long_cnt == NULL) {
2095 lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
2096 if (lz->long_cnt == NULL)
2097 return NULL;
2098 }
2099 if (one == NULL) {
2100 one = PyInt_FromLong(1);
2101 if (one == NULL)
2102 return NULL;
2103 }
2104 cnt = lz->long_cnt;
2105 assert(cnt != NULL);
2106 stepped_up = PyNumber_Add(cnt, one);
2107 if (stepped_up == NULL)
2108 return NULL;
2109 lz->long_cnt = stepped_up;
2110 return cnt;
2111}
2112
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002113static PyObject *
2114count_next(countobject *lz)
2115{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002116 if (lz->cnt == PY_SSIZE_T_MAX)
2117 return count_nextlong(lz);
Jack Diederich36234e82006-09-21 17:50:26 +00002118 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002119}
2120
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002121static PyObject *
2122count_repr(countobject *lz)
2123{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002124 PyObject *cnt_repr;
2125 PyObject *result;
2126
2127 if (lz->cnt != PY_SSIZE_T_MAX)
2128 return PyString_FromFormat("count(%zd)", lz->cnt);
2129
2130 cnt_repr = PyObject_Repr(lz->long_cnt);
2131 if (cnt_repr == NULL)
2132 return NULL;
2133 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
2134 Py_DECREF(cnt_repr);
2135 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002136}
2137
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002138PyDoc_STRVAR(count_doc,
2139"count([firstval]) --> count object\n\
2140\n\
2141Return a count object whose .next() method returns consecutive\n\
2142integers starting from zero or, if specified, from firstval.");
2143
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002144static PyTypeObject count_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002145 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002146 "itertools.count", /* tp_name */
2147 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002148 0, /* tp_itemsize */
2149 /* methods */
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002150 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002151 0, /* tp_print */
2152 0, /* tp_getattr */
2153 0, /* tp_setattr */
2154 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002155 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002156 0, /* tp_as_number */
2157 0, /* tp_as_sequence */
2158 0, /* tp_as_mapping */
2159 0, /* tp_hash */
2160 0, /* tp_call */
2161 0, /* tp_str */
2162 PyObject_GenericGetAttr, /* tp_getattro */
2163 0, /* tp_setattro */
2164 0, /* tp_as_buffer */
2165 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002166 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002167 0, /* tp_traverse */
2168 0, /* tp_clear */
2169 0, /* tp_richcompare */
2170 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002171 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002172 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002173 0, /* tp_methods */
2174 0, /* tp_members */
2175 0, /* tp_getset */
2176 0, /* tp_base */
2177 0, /* tp_dict */
2178 0, /* tp_descr_get */
2179 0, /* tp_descr_set */
2180 0, /* tp_dictoffset */
2181 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002182 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002183 count_new, /* tp_new */
2184};
2185
2186
2187/* izip object ************************************************************/
2188
2189#include "Python.h"
2190
2191typedef struct {
2192 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002193 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002194 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002195 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002196} izipobject;
2197
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002198static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002199
2200static PyObject *
2201izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2202{
2203 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002204 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002205 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002206 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002207 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002208
Georg Brandlb84c1372007-01-21 10:28:43 +00002209 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002210 return NULL;
2211
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002212 /* args must be a tuple */
2213 assert(PyTuple_Check(args));
2214
2215 /* obtain iterators */
2216 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002217 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002218 return NULL;
2219 for (i=0; i < tuplesize; ++i) {
2220 PyObject *item = PyTuple_GET_ITEM(args, i);
2221 PyObject *it = PyObject_GetIter(item);
2222 if (it == NULL) {
2223 if (PyErr_ExceptionMatches(PyExc_TypeError))
2224 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002225 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002226 i+1);
2227 Py_DECREF(ittuple);
2228 return NULL;
2229 }
2230 PyTuple_SET_ITEM(ittuple, i, it);
2231 }
2232
Raymond Hettinger2012f172003-02-07 05:32:58 +00002233 /* create a result holder */
2234 result = PyTuple_New(tuplesize);
2235 if (result == NULL) {
2236 Py_DECREF(ittuple);
2237 return NULL;
2238 }
2239 for (i=0 ; i < tuplesize ; i++) {
2240 Py_INCREF(Py_None);
2241 PyTuple_SET_ITEM(result, i, Py_None);
2242 }
2243
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002244 /* create izipobject structure */
2245 lz = (izipobject *)type->tp_alloc(type, 0);
2246 if (lz == NULL) {
2247 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002248 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002249 return NULL;
2250 }
2251 lz->ittuple = ittuple;
2252 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002253 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002254
2255 return (PyObject *)lz;
2256}
2257
2258static void
2259izip_dealloc(izipobject *lz)
2260{
2261 PyObject_GC_UnTrack(lz);
2262 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002263 Py_XDECREF(lz->result);
Martin v. Löwis68192102007-07-21 06:55:02 +00002264 Py_Type(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002265}
2266
2267static int
2268izip_traverse(izipobject *lz, visitproc visit, void *arg)
2269{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002270 Py_VISIT(lz->ittuple);
2271 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002272 return 0;
2273}
2274
2275static PyObject *
2276izip_next(izipobject *lz)
2277{
Jack Diederich6c433a92006-05-26 11:15:17 +00002278 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002279 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002280 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002281 PyObject *it;
2282 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002283 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002284
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002285 if (tuplesize == 0)
2286 return NULL;
Martin v. Löwis68192102007-07-21 06:55:02 +00002287 if (Py_Refcnt(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002288 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002289 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002290 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002291 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00002292 item = (*Py_Type(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002293 if (item == NULL) {
2294 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002295 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002296 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002297 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002298 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002299 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002300 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002301 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002302 result = PyTuple_New(tuplesize);
2303 if (result == NULL)
2304 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002305 for (i=0 ; i < tuplesize ; i++) {
2306 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002307 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00002308 item = (*Py_Type(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002309 if (item == NULL) {
2310 Py_DECREF(result);
2311 return NULL;
2312 }
2313 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002314 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002315 }
2316 return result;
2317}
2318
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002319PyDoc_STRVAR(izip_doc,
2320"izip(iter1 [,iter2 [...]]) --> izip object\n\
2321\n\
2322Return a izip object whose .next() method returns a tuple where\n\
2323the i-th element comes from the i-th iterable argument. The .next()\n\
2324method continues until the shortest iterable in the argument sequence\n\
2325is exhausted and then it raises StopIteration. Works like the zip()\n\
2326function but consumes less memory by returning an iterator instead of\n\
2327a list.");
2328
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002329static PyTypeObject izip_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002330 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002331 "itertools.izip", /* tp_name */
2332 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002333 0, /* tp_itemsize */
2334 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002335 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002336 0, /* tp_print */
2337 0, /* tp_getattr */
2338 0, /* tp_setattr */
2339 0, /* tp_compare */
2340 0, /* tp_repr */
2341 0, /* tp_as_number */
2342 0, /* tp_as_sequence */
2343 0, /* tp_as_mapping */
2344 0, /* tp_hash */
2345 0, /* tp_call */
2346 0, /* tp_str */
2347 PyObject_GenericGetAttr, /* tp_getattro */
2348 0, /* tp_setattro */
2349 0, /* tp_as_buffer */
2350 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2351 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002352 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002353 (traverseproc)izip_traverse, /* tp_traverse */
2354 0, /* tp_clear */
2355 0, /* tp_richcompare */
2356 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002357 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002358 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002359 0, /* tp_methods */
2360 0, /* tp_members */
2361 0, /* tp_getset */
2362 0, /* tp_base */
2363 0, /* tp_dict */
2364 0, /* tp_descr_get */
2365 0, /* tp_descr_set */
2366 0, /* tp_dictoffset */
2367 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002368 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002369 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002370 PyObject_GC_Del, /* tp_free */
2371};
2372
2373
2374/* repeat object ************************************************************/
2375
2376typedef struct {
2377 PyObject_HEAD
2378 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002379 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002380} repeatobject;
2381
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002382static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002383
2384static PyObject *
2385repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2386{
2387 repeatobject *ro;
2388 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002389 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002390
Georg Brandlb84c1372007-01-21 10:28:43 +00002391 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002392 return NULL;
2393
Jack Diederich6c433a92006-05-26 11:15:17 +00002394 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002395 return NULL;
2396
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002397 if (PyTuple_Size(args) == 2 && cnt < 0)
2398 cnt = 0;
2399
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002400 ro = (repeatobject *)type->tp_alloc(type, 0);
2401 if (ro == NULL)
2402 return NULL;
2403 Py_INCREF(element);
2404 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002405 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002406 return (PyObject *)ro;
2407}
2408
2409static void
2410repeat_dealloc(repeatobject *ro)
2411{
2412 PyObject_GC_UnTrack(ro);
2413 Py_XDECREF(ro->element);
Martin v. Löwis68192102007-07-21 06:55:02 +00002414 Py_Type(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002415}
2416
2417static int
2418repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2419{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002420 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002421 return 0;
2422}
2423
2424static PyObject *
2425repeat_next(repeatobject *ro)
2426{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002427 if (ro->cnt == 0)
2428 return NULL;
2429 if (ro->cnt > 0)
2430 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002431 Py_INCREF(ro->element);
2432 return ro->element;
2433}
2434
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002435static PyObject *
2436repeat_repr(repeatobject *ro)
2437{
2438 PyObject *result, *objrepr;
2439
2440 objrepr = PyObject_Repr(ro->element);
2441 if (objrepr == NULL)
2442 return NULL;
2443
2444 if (ro->cnt == -1)
2445 result = PyString_FromFormat("repeat(%s)",
2446 PyString_AS_STRING(objrepr));
2447 else
Jack Diederich6c433a92006-05-26 11:15:17 +00002448 result = PyString_FromFormat("repeat(%s, %zd)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002449 PyString_AS_STRING(objrepr), ro->cnt);
2450 Py_DECREF(objrepr);
2451 return result;
2452}
2453
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002454static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002455repeat_len(repeatobject *ro)
2456{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002457 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002458 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002459 return NULL;
2460 }
Jack Diederich6c433a92006-05-26 11:15:17 +00002461 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002462}
2463
Armin Rigof5b3e362006-02-11 21:32:43 +00002464PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002465
2466static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002467 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002468 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002469};
2470
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002471PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002472"repeat(element [,times]) -> create an iterator which returns the element\n\
2473for the specified number of times. If not specified, returns the element\n\
2474endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002475
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002476static PyTypeObject repeat_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002477 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002478 "itertools.repeat", /* tp_name */
2479 sizeof(repeatobject), /* tp_basicsize */
2480 0, /* tp_itemsize */
2481 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002482 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002483 0, /* tp_print */
2484 0, /* tp_getattr */
2485 0, /* tp_setattr */
2486 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002487 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002488 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002489 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002490 0, /* tp_as_mapping */
2491 0, /* tp_hash */
2492 0, /* tp_call */
2493 0, /* tp_str */
2494 PyObject_GenericGetAttr, /* tp_getattro */
2495 0, /* tp_setattro */
2496 0, /* tp_as_buffer */
2497 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2498 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002499 repeat_doc, /* tp_doc */
2500 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002501 0, /* tp_clear */
2502 0, /* tp_richcompare */
2503 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002504 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002505 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002506 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002507 0, /* tp_members */
2508 0, /* tp_getset */
2509 0, /* tp_base */
2510 0, /* tp_dict */
2511 0, /* tp_descr_get */
2512 0, /* tp_descr_set */
2513 0, /* tp_dictoffset */
2514 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002515 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002516 repeat_new, /* tp_new */
2517 PyObject_GC_Del, /* tp_free */
2518};
2519
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002520/* iziplongest object ************************************************************/
2521
2522#include "Python.h"
2523
2524typedef struct {
2525 PyObject_HEAD
2526 Py_ssize_t tuplesize;
2527 Py_ssize_t numactive;
2528 PyObject *ittuple; /* tuple of iterators */
2529 PyObject *result;
2530 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002531} iziplongestobject;
2532
2533static PyTypeObject iziplongest_type;
2534
2535static PyObject *
2536izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2537{
2538 iziplongestobject *lz;
2539 Py_ssize_t i;
2540 PyObject *ittuple; /* tuple of iterators */
2541 PyObject *result;
2542 PyObject *fillvalue = Py_None;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002543 Py_ssize_t tuplesize = PySequence_Length(args);
2544
2545 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
2546 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
2547 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
2548 PyErr_SetString(PyExc_TypeError,
2549 "izip_longest() got an unexpected keyword argument");
2550 return NULL;
2551 }
2552 }
2553
2554 /* args must be a tuple */
2555 assert(PyTuple_Check(args));
2556
2557 /* obtain iterators */
2558 ittuple = PyTuple_New(tuplesize);
2559 if (ittuple == NULL)
2560 return NULL;
2561 for (i=0; i < tuplesize; ++i) {
2562 PyObject *item = PyTuple_GET_ITEM(args, i);
2563 PyObject *it = PyObject_GetIter(item);
2564 if (it == NULL) {
2565 if (PyErr_ExceptionMatches(PyExc_TypeError))
2566 PyErr_Format(PyExc_TypeError,
2567 "izip_longest argument #%zd must support iteration",
2568 i+1);
2569 Py_DECREF(ittuple);
2570 return NULL;
2571 }
2572 PyTuple_SET_ITEM(ittuple, i, it);
2573 }
2574
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002575 /* create a result holder */
2576 result = PyTuple_New(tuplesize);
2577 if (result == NULL) {
2578 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002579 return NULL;
2580 }
2581 for (i=0 ; i < tuplesize ; i++) {
2582 Py_INCREF(Py_None);
2583 PyTuple_SET_ITEM(result, i, Py_None);
2584 }
2585
2586 /* create iziplongestobject structure */
2587 lz = (iziplongestobject *)type->tp_alloc(type, 0);
2588 if (lz == NULL) {
2589 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002590 Py_DECREF(result);
2591 return NULL;
2592 }
2593 lz->ittuple = ittuple;
2594 lz->tuplesize = tuplesize;
2595 lz->numactive = tuplesize;
2596 lz->result = result;
2597 Py_INCREF(fillvalue);
2598 lz->fillvalue = fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002599 return (PyObject *)lz;
2600}
2601
2602static void
2603izip_longest_dealloc(iziplongestobject *lz)
2604{
2605 PyObject_GC_UnTrack(lz);
2606 Py_XDECREF(lz->ittuple);
2607 Py_XDECREF(lz->result);
2608 Py_XDECREF(lz->fillvalue);
Martin v. Löwis68192102007-07-21 06:55:02 +00002609 Py_Type(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002610}
2611
2612static int
2613izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
2614{
2615 Py_VISIT(lz->ittuple);
2616 Py_VISIT(lz->result);
2617 Py_VISIT(lz->fillvalue);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002618 return 0;
2619}
2620
2621static PyObject *
2622izip_longest_next(iziplongestobject *lz)
2623{
2624 Py_ssize_t i;
2625 Py_ssize_t tuplesize = lz->tuplesize;
2626 PyObject *result = lz->result;
2627 PyObject *it;
2628 PyObject *item;
2629 PyObject *olditem;
2630
2631 if (tuplesize == 0)
2632 return NULL;
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002633 if (lz->numactive == 0)
2634 return NULL;
Martin v. Löwis68192102007-07-21 06:55:02 +00002635 if (Py_Refcnt(result) == 1) {
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002636 Py_INCREF(result);
2637 for (i=0 ; i < tuplesize ; i++) {
2638 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002639 if (it == NULL) {
2640 Py_INCREF(lz->fillvalue);
2641 item = lz->fillvalue;
2642 } else {
2643 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00002644 item = (*Py_Type(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002645 if (item == NULL) {
2646 lz->numactive -= 1;
2647 if (lz->numactive == 0) {
2648 Py_DECREF(result);
2649 return NULL;
2650 } else {
2651 Py_INCREF(lz->fillvalue);
2652 item = lz->fillvalue;
2653 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
2654 Py_DECREF(it);
2655 }
2656 }
2657 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002658 olditem = PyTuple_GET_ITEM(result, i);
2659 PyTuple_SET_ITEM(result, i, item);
2660 Py_DECREF(olditem);
2661 }
2662 } else {
2663 result = PyTuple_New(tuplesize);
2664 if (result == NULL)
2665 return NULL;
2666 for (i=0 ; i < tuplesize ; i++) {
2667 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002668 if (it == NULL) {
2669 Py_INCREF(lz->fillvalue);
2670 item = lz->fillvalue;
2671 } else {
2672 assert(PyIter_Check(it));
Martin v. Löwis68192102007-07-21 06:55:02 +00002673 item = (*Py_Type(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002674 if (item == NULL) {
2675 lz->numactive -= 1;
2676 if (lz->numactive == 0) {
2677 Py_DECREF(result);
2678 return NULL;
2679 } else {
2680 Py_INCREF(lz->fillvalue);
2681 item = lz->fillvalue;
2682 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
2683 Py_DECREF(it);
2684 }
2685 }
2686 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002687 PyTuple_SET_ITEM(result, i, item);
2688 }
2689 }
2690 return result;
2691}
2692
2693PyDoc_STRVAR(izip_longest_doc,
2694"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
2695\n\
2696Return an izip_longest object whose .next() method returns a tuple where\n\
2697the i-th element comes from the i-th iterable argument. The .next()\n\
2698method continues until the longest iterable in the argument sequence\n\
2699is exhausted and then it raises StopIteration. When the shorter iterables\n\
2700are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
2701defaults to None or can be specified by a keyword argument.\n\
2702");
2703
2704static PyTypeObject iziplongest_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002705 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002706 "itertools.izip_longest", /* tp_name */
2707 sizeof(iziplongestobject), /* tp_basicsize */
2708 0, /* tp_itemsize */
2709 /* methods */
2710 (destructor)izip_longest_dealloc, /* tp_dealloc */
2711 0, /* tp_print */
2712 0, /* tp_getattr */
2713 0, /* tp_setattr */
2714 0, /* tp_compare */
2715 0, /* tp_repr */
2716 0, /* tp_as_number */
2717 0, /* tp_as_sequence */
2718 0, /* tp_as_mapping */
2719 0, /* tp_hash */
2720 0, /* tp_call */
2721 0, /* tp_str */
2722 PyObject_GenericGetAttr, /* tp_getattro */
2723 0, /* tp_setattro */
2724 0, /* tp_as_buffer */
2725 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2726 Py_TPFLAGS_BASETYPE, /* tp_flags */
2727 izip_longest_doc, /* tp_doc */
2728 (traverseproc)izip_longest_traverse, /* tp_traverse */
2729 0, /* tp_clear */
2730 0, /* tp_richcompare */
2731 0, /* tp_weaklistoffset */
2732 PyObject_SelfIter, /* tp_iter */
2733 (iternextfunc)izip_longest_next, /* tp_iternext */
2734 0, /* tp_methods */
2735 0, /* tp_members */
2736 0, /* tp_getset */
2737 0, /* tp_base */
2738 0, /* tp_dict */
2739 0, /* tp_descr_get */
2740 0, /* tp_descr_set */
2741 0, /* tp_dictoffset */
2742 0, /* tp_init */
2743 0, /* tp_alloc */
2744 izip_longest_new, /* tp_new */
2745 PyObject_GC_Del, /* tp_free */
2746};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002747
2748/* module level code ********************************************************/
2749
2750PyDoc_STRVAR(module_doc,
2751"Functional tools for creating and using iterators.\n\
2752\n\
2753Infinite iterators:\n\
2754count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002755cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002756repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002757\n\
2758Iterators terminating on the shortest input sequence:\n\
2759izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002760izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002761ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2762ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002763islice(seq, [start,] stop [, step]) --> elements from\n\
2764 seq[start:stop:step]\n\
2765imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2766starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002767tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002768chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002769takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2770dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002771groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002772");
2773
2774
Raymond Hettingerad983e72003-11-12 14:32:26 +00002775static PyMethodDef module_methods[] = {
2776 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2777 {NULL, NULL} /* sentinel */
2778};
2779
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002780PyMODINIT_FUNC
2781inititertools(void)
2782{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002783 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002784 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002785 char *name;
2786 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002787 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002788 &dropwhile_type,
2789 &takewhile_type,
2790 &islice_type,
2791 &starmap_type,
2792 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002793 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002794 &ifilter_type,
2795 &ifilterfalse_type,
2796 &count_type,
2797 &izip_type,
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002798 &iziplongest_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002799 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002800 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002801 NULL
2802 };
2803
Martin v. Löwis68192102007-07-21 06:55:02 +00002804 Py_Type(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00002805 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002806 if (m == NULL)
2807 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002808
Raymond Hettinger60eca932003-02-09 06:40:58 +00002809 for (i=0 ; typelist[i] != NULL ; i++) {
2810 if (PyType_Ready(typelist[i]) < 0)
2811 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002812 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002813 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002814 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00002815 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002816 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00002817
2818 if (PyType_Ready(&teedataobject_type) < 0)
2819 return;
2820 if (PyType_Ready(&tee_type) < 0)
2821 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002822 if (PyType_Ready(&_grouper_type) < 0)
2823 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002824}