blob: 89293098d508945eb22460b59d7ce311a44508a1 [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);
Christian Heimese93237d2007-12-19 02:37:44 +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);
Christian Heimese93237d2007-12-19 02:37:44 +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);
Christian Heimese93237d2007-12-19 02:37:44 +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));
Christian Heimese93237d2007-12-19 02:37:44 +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);
Christian Heimese93237d2007-12-19 02:37:44 +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));
Christian Heimese93237d2007-12-19 02:37:44 +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);
Christian Heimese93237d2007-12-19 02:37:44 +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));
Christian Heimese93237d2007-12-19 02:37:44 +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);
Christian Heimese93237d2007-12-19 02:37:44 +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));
Christian Heimese93237d2007-12-19 02:37:44 +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)) {
Raymond Hettinger47317092008-01-17 03:02:14 +00001359 PyObject *newargs = PySequence_Tuple(args);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001360 Py_DECREF(args);
Raymond Hettinger47317092008-01-17 03:02:14 +00001361 if (newargs == NULL)
1362 return NULL;
1363 args = newargs;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001364 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001365 result = PyObject_Call(lz->func, args, NULL);
1366 Py_DECREF(args);
1367 return result;
1368}
1369
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370PyDoc_STRVAR(starmap_doc,
1371"starmap(function, sequence) --> starmap object\n\
1372\n\
1373Return an iterator whose values are returned from the function evaluated\n\
1374with a argument tuple taken from the given sequence.");
1375
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001376static PyTypeObject starmap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001377 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001378 "itertools.starmap", /* tp_name */
1379 sizeof(starmapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380 0, /* tp_itemsize */
1381 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001382 (destructor)starmap_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383 0, /* tp_print */
1384 0, /* tp_getattr */
1385 0, /* tp_setattr */
1386 0, /* tp_compare */
1387 0, /* tp_repr */
1388 0, /* tp_as_number */
1389 0, /* tp_as_sequence */
1390 0, /* tp_as_mapping */
1391 0, /* tp_hash */
1392 0, /* tp_call */
1393 0, /* tp_str */
1394 PyObject_GenericGetAttr, /* tp_getattro */
1395 0, /* tp_setattro */
1396 0, /* tp_as_buffer */
1397 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1398 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001399 starmap_doc, /* tp_doc */
1400 (traverseproc)starmap_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001401 0, /* tp_clear */
1402 0, /* tp_richcompare */
1403 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001404 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001405 (iternextfunc)starmap_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001406 0, /* tp_methods */
1407 0, /* tp_members */
1408 0, /* tp_getset */
1409 0, /* tp_base */
1410 0, /* tp_dict */
1411 0, /* tp_descr_get */
1412 0, /* tp_descr_set */
1413 0, /* tp_dictoffset */
1414 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001415 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001416 starmap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001417 PyObject_GC_Del, /* tp_free */
1418};
1419
1420
1421/* imap object ************************************************************/
1422
1423typedef struct {
1424 PyObject_HEAD
1425 PyObject *iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001426 PyObject *func;
1427} imapobject;
1428
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001429static PyTypeObject imap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430
1431static PyObject *
1432imap_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1433{
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001434 PyObject *it, *iters, *func;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001435 imapobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001436 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001437
Georg Brandlb84c1372007-01-21 10:28:43 +00001438 if (type == &imap_type && !_PyArg_NoKeywords("imap()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001439 return NULL;
1440
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001441 numargs = PyTuple_Size(args);
1442 if (numargs < 2) {
1443 PyErr_SetString(PyExc_TypeError,
1444 "imap() must have at least two arguments.");
1445 return NULL;
1446 }
1447
1448 iters = PyTuple_New(numargs-1);
1449 if (iters == NULL)
1450 return NULL;
1451
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001452 for (i=1 ; i<numargs ; i++) {
1453 /* Get iterator. */
1454 it = PyObject_GetIter(PyTuple_GET_ITEM(args, i));
1455 if (it == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456 Py_DECREF(iters);
1457 return NULL;
1458 }
1459 PyTuple_SET_ITEM(iters, i-1, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460 }
1461
1462 /* create imapobject structure */
1463 lz = (imapobject *)type->tp_alloc(type, 0);
1464 if (lz == NULL) {
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001465 Py_DECREF(iters);
1466 return NULL;
1467 }
1468 lz->iters = iters;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469 func = PyTuple_GET_ITEM(args, 0);
1470 Py_INCREF(func);
1471 lz->func = func;
1472
1473 return (PyObject *)lz;
1474}
1475
1476static void
1477imap_dealloc(imapobject *lz)
1478{
1479 PyObject_GC_UnTrack(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480 Py_XDECREF(lz->iters);
1481 Py_XDECREF(lz->func);
Christian Heimese93237d2007-12-19 02:37:44 +00001482 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001483}
1484
1485static int
1486imap_traverse(imapobject *lz, visitproc visit, void *arg)
1487{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001488 Py_VISIT(lz->iters);
1489 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490 return 0;
1491}
1492
Raymond Hettinger2012f172003-02-07 05:32:58 +00001493/*
1494imap() is an iterator version of __builtins__.map() except that it does
1495not have the None fill-in feature. That was intentionally left out for
1496the following reasons:
1497
1498 1) Itertools are designed to be easily combined and chained together.
1499 Having all tools stop with the shortest input is a unifying principle
1500 that makes it easier to combine finite iterators (supplying data) with
1501 infinite iterators like count() and repeat() (for supplying sequential
1502 or constant arguments to a function).
1503
1504 2) In typical use cases for combining itertools, having one finite data
1505 supplier run out before another is likely to be an error condition which
1506 should not pass silently by automatically supplying None.
1507
1508 3) The use cases for automatic None fill-in are rare -- not many functions
1509 do something useful when a parameter suddenly switches type and becomes
1510 None.
1511
1512 4) If a need does arise, it can be met by __builtins__.map() or by
Raymond Hettingerbefa37d2003-06-18 19:25:37 +00001513 writing: chain(iterable, repeat(None)).
Raymond Hettinger2012f172003-02-07 05:32:58 +00001514
1515 5) Similar toolsets in Haskell and SML do not have automatic None fill-in.
1516*/
1517
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001518static PyObject *
1519imap_next(imapobject *lz)
1520{
1521 PyObject *val;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001522 PyObject *argtuple;
1523 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001524 Py_ssize_t numargs, i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001525
1526 numargs = PyTuple_Size(lz->iters);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001527 argtuple = PyTuple_New(numargs);
1528 if (argtuple == NULL)
1529 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001530
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001531 for (i=0 ; i<numargs ; i++) {
1532 val = PyIter_Next(PyTuple_GET_ITEM(lz->iters, i));
1533 if (val == NULL) {
1534 Py_DECREF(argtuple);
1535 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001537 PyTuple_SET_ITEM(argtuple, i, val);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001538 }
Raymond Hettingerf0c00242003-02-07 07:26:25 +00001539 if (lz->func == Py_None)
1540 return argtuple;
1541 result = PyObject_Call(lz->func, argtuple, NULL);
1542 Py_DECREF(argtuple);
1543 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001544}
1545
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001546PyDoc_STRVAR(imap_doc,
1547"imap(func, *iterables) --> imap object\n\
1548\n\
1549Make an iterator that computes the function using arguments from\n\
1550each of the iterables. Like map() except that it returns\n\
1551an iterator instead of a list and that it stops when the shortest\n\
1552iterable is exhausted instead of filling in None for shorter\n\
1553iterables.");
1554
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001555static PyTypeObject imap_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001556 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00001557 "itertools.imap", /* tp_name */
1558 sizeof(imapobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001559 0, /* tp_itemsize */
1560 /* methods */
1561 (destructor)imap_dealloc, /* tp_dealloc */
1562 0, /* tp_print */
1563 0, /* tp_getattr */
1564 0, /* tp_setattr */
1565 0, /* tp_compare */
1566 0, /* tp_repr */
1567 0, /* tp_as_number */
1568 0, /* tp_as_sequence */
1569 0, /* tp_as_mapping */
1570 0, /* tp_hash */
1571 0, /* tp_call */
1572 0, /* tp_str */
1573 PyObject_GenericGetAttr, /* tp_getattro */
1574 0, /* tp_setattro */
1575 0, /* tp_as_buffer */
1576 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1577 Py_TPFLAGS_BASETYPE, /* tp_flags */
1578 imap_doc, /* tp_doc */
1579 (traverseproc)imap_traverse, /* tp_traverse */
1580 0, /* tp_clear */
1581 0, /* tp_richcompare */
1582 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001583 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001584 (iternextfunc)imap_next, /* tp_iternext */
1585 0, /* tp_methods */
1586 0, /* tp_members */
1587 0, /* tp_getset */
1588 0, /* tp_base */
1589 0, /* tp_dict */
1590 0, /* tp_descr_get */
1591 0, /* tp_descr_set */
1592 0, /* tp_dictoffset */
1593 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001594 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00001595 imap_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001596 PyObject_GC_Del, /* tp_free */
1597};
1598
1599
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001600/* chain object ************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001601
1602typedef struct {
1603 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00001604 Py_ssize_t tuplesize;
1605 Py_ssize_t iternum; /* which iterator is active */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001606 PyObject *ittuple; /* tuple of iterators */
1607} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001609static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001610
1611static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001612chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001613{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001614 chainobject *lz;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00001615 Py_ssize_t tuplesize = PySequence_Length(args);
Jack Diederich6c433a92006-05-26 11:15:17 +00001616 Py_ssize_t i;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001617 PyObject *ittuple;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618
Georg Brandlb84c1372007-01-21 10:28:43 +00001619 if (type == &chain_type && !_PyArg_NoKeywords("chain()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001620 return NULL;
1621
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001622 /* obtain iterators */
1623 assert(PyTuple_Check(args));
1624 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001625 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001626 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001627 for (i=0; i < tuplesize; ++i) {
1628 PyObject *item = PyTuple_GET_ITEM(args, i);
1629 PyObject *it = PyObject_GetIter(item);
1630 if (it == NULL) {
1631 if (PyErr_ExceptionMatches(PyExc_TypeError))
1632 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00001633 "chain argument #%zd must support iteration",
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001634 i+1);
1635 Py_DECREF(ittuple);
1636 return NULL;
1637 }
1638 PyTuple_SET_ITEM(ittuple, i, it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001639 }
1640
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001641 /* create chainobject structure */
1642 lz = (chainobject *)type->tp_alloc(type, 0);
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001643 if (lz == NULL) {
1644 Py_DECREF(ittuple);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001645 return NULL;
Raymond Hettinger7d98fb92003-06-17 23:14:40 +00001646 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001647
1648 lz->ittuple = ittuple;
1649 lz->iternum = 0;
1650 lz->tuplesize = tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001651
1652 return (PyObject *)lz;
1653}
1654
1655static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001656chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001657{
1658 PyObject_GC_UnTrack(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001659 Py_XDECREF(lz->ittuple);
Christian Heimese93237d2007-12-19 02:37:44 +00001660 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001661}
1662
Raymond Hettinger2012f172003-02-07 05:32:58 +00001663static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001664chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001665{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00001666 Py_VISIT(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00001667 return 0;
1668}
1669
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001671chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001672{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001673 PyObject *it;
1674 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001676 while (lz->iternum < lz->tuplesize) {
1677 it = PyTuple_GET_ITEM(lz->ittuple, lz->iternum);
1678 item = PyIter_Next(it);
1679 if (item != NULL)
1680 return item;
Raymond Hettinger9d7c8702004-05-08 19:49:42 +00001681 if (PyErr_Occurred()) {
1682 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1683 PyErr_Clear();
1684 else
1685 return NULL;
1686 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001687 lz->iternum++;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001688 }
1689 return NULL;
1690}
1691
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001692PyDoc_STRVAR(chain_doc,
1693"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001694\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001695Return a chain object whose .next() method returns elements from the\n\
1696first iterable until it is exhausted, then elements from the next\n\
1697iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001698
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001699static PyTypeObject chain_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00001700 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001701 "itertools.chain", /* tp_name */
1702 sizeof(chainobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001703 0, /* tp_itemsize */
1704 /* methods */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001705 (destructor)chain_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001706 0, /* tp_print */
1707 0, /* tp_getattr */
1708 0, /* tp_setattr */
1709 0, /* tp_compare */
1710 0, /* tp_repr */
1711 0, /* tp_as_number */
1712 0, /* tp_as_sequence */
1713 0, /* tp_as_mapping */
1714 0, /* tp_hash */
1715 0, /* tp_call */
1716 0, /* tp_str */
1717 PyObject_GenericGetAttr, /* tp_getattro */
1718 0, /* tp_setattro */
1719 0, /* tp_as_buffer */
1720 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1721 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001722 chain_doc, /* tp_doc */
1723 (traverseproc)chain_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724 0, /* tp_clear */
1725 0, /* tp_richcompare */
1726 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00001727 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001728 (iternextfunc)chain_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001729 0, /* tp_methods */
1730 0, /* tp_members */
1731 0, /* tp_getset */
1732 0, /* tp_base */
1733 0, /* tp_dict */
1734 0, /* tp_descr_get */
1735 0, /* tp_descr_set */
1736 0, /* tp_dictoffset */
1737 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00001738 0, /* tp_alloc */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001739 chain_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001740 PyObject_GC_Del, /* tp_free */
1741};
1742
1743
Raymond Hettinger50986cc2008-02-22 03:16:42 +00001744/* product object ************************************************************/
1745
1746typedef struct {
1747 PyObject_HEAD
1748 PyObject *pools; /* tuple of pool tuples */
1749 Py_ssize_t *maxvec;
1750 Py_ssize_t *indices;
1751 PyObject *result;
1752 int stopped;
1753} productobject;
1754
1755static PyTypeObject product_type;
1756
1757static PyObject *
1758product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1759{
1760 productobject *lz;
1761 Py_ssize_t npools;
1762 PyObject *pools = NULL;
1763 Py_ssize_t *maxvec = NULL;
1764 Py_ssize_t *indices = NULL;
1765 Py_ssize_t i;
1766
1767 if (type == &product_type && !_PyArg_NoKeywords("product()", kwds))
1768 return NULL;
1769
1770 assert(PyTuple_Check(args));
1771 npools = PyTuple_GET_SIZE(args);
1772
1773 maxvec = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1774 indices = PyMem_Malloc(npools * sizeof(Py_ssize_t));
1775 if (maxvec == NULL || indices == NULL) {
1776 PyErr_NoMemory();
1777 goto error;
1778 }
1779
1780 pools = PyTuple_New(npools);
1781 if (pools == NULL)
1782 goto error;
1783
1784 for (i=0; i < npools; ++i) {
1785 PyObject *item = PyTuple_GET_ITEM(args, i);
1786 PyObject *pool = PySequence_Tuple(item);
1787 if (pool == NULL)
1788 goto error;
1789
1790 PyTuple_SET_ITEM(pools, i, pool);
1791 maxvec[i] = PyTuple_GET_SIZE(pool);
1792 indices[i] = 0;
1793 }
1794
1795 /* create productobject structure */
1796 lz = (productobject *)type->tp_alloc(type, 0);
1797 if (lz == NULL) {
1798 Py_DECREF(pools);
1799 return NULL;
1800 }
1801
1802 lz->pools = pools;
1803 lz->maxvec = maxvec;
1804 lz->indices = indices;
1805 lz->result = NULL;
1806 lz->stopped = 0;
1807
1808 return (PyObject *)lz;
1809
1810error:
1811 if (maxvec != NULL)
1812 PyMem_Free(maxvec);
1813 if (indices != NULL)
1814 PyMem_Free(indices);
1815 Py_XDECREF(pools);
1816 return NULL;
1817}
1818
1819static void
1820product_dealloc(productobject *lz)
1821{
1822 PyObject_GC_UnTrack(lz);
1823 Py_XDECREF(lz->pools);
1824 Py_XDECREF(lz->result);
1825 PyMem_Free(lz->maxvec);
1826 PyMem_Free(lz->indices);
1827 Py_TYPE(lz)->tp_free(lz);
1828}
1829
1830static int
1831product_traverse(productobject *lz, visitproc visit, void *arg)
1832{
1833 Py_VISIT(lz->pools);
1834 Py_VISIT(lz->result);
1835 return 0;
1836}
1837
1838static PyObject *
1839product_next(productobject *lz)
1840{
1841 PyObject *pool;
1842 PyObject *elem;
1843 PyObject *tuple_result;
1844 PyObject *pools = lz->pools;
1845 PyObject *result = lz->result;
1846 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
1847 Py_ssize_t i;
1848
1849 if (lz->stopped)
1850 return NULL;
1851 if (result == NULL) {
1852 if (npools == 0)
1853 goto empty;
1854 result = PyList_New(npools);
1855 if (result == NULL)
1856 goto empty;
1857 lz->result = result;
1858 for (i=0; i < npools; i++) {
1859 pool = PyTuple_GET_ITEM(pools, i);
1860 if (PyTuple_GET_SIZE(pool) == 0)
1861 goto empty;
1862 elem = PyTuple_GET_ITEM(pool, 0);
1863 Py_INCREF(elem);
1864 PyList_SET_ITEM(result, i, elem);
1865 }
1866 } else {
1867 Py_ssize_t *indices = lz->indices;
1868 Py_ssize_t *maxvec = lz->maxvec;
1869 for (i=npools-1 ; i >= 0 ; i--) {
1870 pool = PyTuple_GET_ITEM(pools, i);
1871 indices[i]++;
1872 if (indices[i] == maxvec[i]) {
1873 indices[i] = 0;
1874 elem = PyTuple_GET_ITEM(pool, 0);
1875 Py_INCREF(elem);
1876 PyList_SetItem(result, i, elem);
1877 } else {
1878 elem = PyTuple_GET_ITEM(pool, indices[i]);
1879 Py_INCREF(elem);
1880 PyList_SetItem(result, i, elem);
1881 break;
1882 }
1883 }
1884 if (i < 0)
1885 return NULL;
1886 }
1887
1888 tuple_result = PySequence_Tuple(result);
1889 if (tuple_result == NULL)
1890 lz->stopped = 1;
1891 return tuple_result;
1892
1893empty:
1894 lz->stopped = 1;
1895 return NULL;
1896}
1897
1898PyDoc_STRVAR(product_doc,
1899"product(*iterables) --> product object\n\
1900\n\
1901Cartesian product of input interables. Equivalent to nested for-loops.\n\n\
1902For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
1903The leftmost iterators are in the outermost for-loop, so the output tuples\n\
1904cycle in a manner similar to an odometer (with the rightmost element changing\n\
1905on every iteration).\n\n\
1906product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
1907product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
1908
1909static PyTypeObject product_type = {
1910 PyVarObject_HEAD_INIT(NULL, 0)
1911 "itertools.product", /* tp_name */
1912 sizeof(productobject), /* tp_basicsize */
1913 0, /* tp_itemsize */
1914 /* methods */
1915 (destructor)product_dealloc, /* tp_dealloc */
1916 0, /* tp_print */
1917 0, /* tp_getattr */
1918 0, /* tp_setattr */
1919 0, /* tp_compare */
1920 0, /* tp_repr */
1921 0, /* tp_as_number */
1922 0, /* tp_as_sequence */
1923 0, /* tp_as_mapping */
1924 0, /* tp_hash */
1925 0, /* tp_call */
1926 0, /* tp_str */
1927 PyObject_GenericGetAttr, /* tp_getattro */
1928 0, /* tp_setattro */
1929 0, /* tp_as_buffer */
1930 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1931 Py_TPFLAGS_BASETYPE, /* tp_flags */
1932 product_doc, /* tp_doc */
1933 (traverseproc)product_traverse, /* tp_traverse */
1934 0, /* tp_clear */
1935 0, /* tp_richcompare */
1936 0, /* tp_weaklistoffset */
1937 PyObject_SelfIter, /* tp_iter */
1938 (iternextfunc)product_next, /* tp_iternext */
1939 0, /* tp_methods */
1940 0, /* tp_members */
1941 0, /* tp_getset */
1942 0, /* tp_base */
1943 0, /* tp_dict */
1944 0, /* tp_descr_get */
1945 0, /* tp_descr_set */
1946 0, /* tp_dictoffset */
1947 0, /* tp_init */
1948 0, /* tp_alloc */
1949 product_new, /* tp_new */
1950 PyObject_GC_Del, /* tp_free */
1951};
1952
1953
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001954/* ifilter object ************************************************************/
1955
1956typedef struct {
1957 PyObject_HEAD
1958 PyObject *func;
1959 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001960} ifilterobject;
1961
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001962static PyTypeObject ifilter_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001963
1964static PyObject *
1965ifilter_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1966{
Raymond Hettinger60eca932003-02-09 06:40:58 +00001967 PyObject *func, *seq;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001968 PyObject *it;
1969 ifilterobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001970
Georg Brandlb84c1372007-01-21 10:28:43 +00001971 if (type == &ifilter_type && !_PyArg_NoKeywords("ifilter()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00001972 return NULL;
1973
Raymond Hettinger60eca932003-02-09 06:40:58 +00001974 if (!PyArg_UnpackTuple(args, "ifilter", 2, 2, &func, &seq))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001975 return NULL;
1976
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001977 /* Get iterator. */
1978 it = PyObject_GetIter(seq);
1979 if (it == NULL)
1980 return NULL;
1981
1982 /* create ifilterobject structure */
1983 lz = (ifilterobject *)type->tp_alloc(type, 0);
1984 if (lz == NULL) {
1985 Py_DECREF(it);
1986 return NULL;
1987 }
1988 Py_INCREF(func);
1989 lz->func = func;
1990 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001991
1992 return (PyObject *)lz;
1993}
1994
1995static void
1996ifilter_dealloc(ifilterobject *lz)
1997{
1998 PyObject_GC_UnTrack(lz);
1999 Py_XDECREF(lz->func);
2000 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002001 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002002}
2003
2004static int
2005ifilter_traverse(ifilterobject *lz, visitproc visit, void *arg)
2006{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002007 Py_VISIT(lz->it);
2008 Py_VISIT(lz->func);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002009 return 0;
2010}
2011
2012static PyObject *
2013ifilter_next(ifilterobject *lz)
2014{
2015 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002016 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002017 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002018 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002019
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002020 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002021 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002022 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002023 item = iternext(it);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002024 if (item == NULL)
2025 return NULL;
2026
2027 if (lz->func == Py_None) {
2028 ok = PyObject_IsTrue(item);
2029 } else {
2030 PyObject *good;
2031 good = PyObject_CallFunctionObjArgs(lz->func,
2032 item, NULL);
2033 if (good == NULL) {
2034 Py_DECREF(item);
2035 return NULL;
2036 }
2037 ok = PyObject_IsTrue(good);
2038 Py_DECREF(good);
2039 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00002040 if (ok)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002041 return item;
2042 Py_DECREF(item);
2043 }
2044}
2045
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002046PyDoc_STRVAR(ifilter_doc,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002047"ifilter(function or None, sequence) --> ifilter object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002048\n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002049Return those items of sequence for which function(item) is true.\n\
2050If function is None, return the items that are true.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002051
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002052static PyTypeObject ifilter_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002053 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002054 "itertools.ifilter", /* tp_name */
2055 sizeof(ifilterobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002056 0, /* tp_itemsize */
2057 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002058 (destructor)ifilter_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002059 0, /* tp_print */
2060 0, /* tp_getattr */
2061 0, /* tp_setattr */
2062 0, /* tp_compare */
2063 0, /* tp_repr */
2064 0, /* tp_as_number */
2065 0, /* tp_as_sequence */
2066 0, /* tp_as_mapping */
2067 0, /* tp_hash */
2068 0, /* tp_call */
2069 0, /* tp_str */
2070 PyObject_GenericGetAttr, /* tp_getattro */
2071 0, /* tp_setattro */
2072 0, /* tp_as_buffer */
2073 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2074 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002075 ifilter_doc, /* tp_doc */
2076 (traverseproc)ifilter_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002077 0, /* tp_clear */
2078 0, /* tp_richcompare */
2079 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002080 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002081 (iternextfunc)ifilter_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002082 0, /* tp_methods */
2083 0, /* tp_members */
2084 0, /* tp_getset */
2085 0, /* tp_base */
2086 0, /* tp_dict */
2087 0, /* tp_descr_get */
2088 0, /* tp_descr_set */
2089 0, /* tp_dictoffset */
2090 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002091 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002092 ifilter_new, /* tp_new */
2093 PyObject_GC_Del, /* tp_free */
2094};
2095
2096
2097/* ifilterfalse object ************************************************************/
2098
2099typedef struct {
2100 PyObject_HEAD
2101 PyObject *func;
2102 PyObject *it;
2103} ifilterfalseobject;
2104
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002105static PyTypeObject ifilterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002106
2107static PyObject *
2108ifilterfalse_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2109{
Guido van Rossumd58f3fc2003-02-09 17:19:18 +00002110 PyObject *func, *seq;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002111 PyObject *it;
2112 ifilterfalseobject *lz;
2113
Georg Brandlb84c1372007-01-21 10:28:43 +00002114 if (type == &ifilterfalse_type &&
2115 !_PyArg_NoKeywords("ifilterfalse()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002116 return NULL;
2117
Raymond Hettinger60eca932003-02-09 06:40:58 +00002118 if (!PyArg_UnpackTuple(args, "ifilterfalse", 2, 2, &func, &seq))
2119 return NULL;
2120
2121 /* Get iterator. */
2122 it = PyObject_GetIter(seq);
2123 if (it == NULL)
2124 return NULL;
2125
2126 /* create ifilterfalseobject structure */
2127 lz = (ifilterfalseobject *)type->tp_alloc(type, 0);
2128 if (lz == NULL) {
2129 Py_DECREF(it);
2130 return NULL;
2131 }
2132 Py_INCREF(func);
2133 lz->func = func;
2134 lz->it = it;
2135
2136 return (PyObject *)lz;
2137}
2138
2139static void
2140ifilterfalse_dealloc(ifilterfalseobject *lz)
2141{
2142 PyObject_GC_UnTrack(lz);
2143 Py_XDECREF(lz->func);
2144 Py_XDECREF(lz->it);
Christian Heimese93237d2007-12-19 02:37:44 +00002145 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002146}
2147
2148static int
2149ifilterfalse_traverse(ifilterfalseobject *lz, visitproc visit, void *arg)
2150{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002151 Py_VISIT(lz->it);
2152 Py_VISIT(lz->func);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002153 return 0;
2154}
2155
2156static PyObject *
2157ifilterfalse_next(ifilterfalseobject *lz)
2158{
2159 PyObject *item;
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002160 PyObject *it = lz->it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002161 long ok;
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002162 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002163
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002164 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002165 iternext = *Py_TYPE(it)->tp_iternext;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002166 for (;;) {
Raymond Hettinger0faa1ca2004-03-17 04:27:44 +00002167 item = iternext(it);
Raymond Hettinger60eca932003-02-09 06:40:58 +00002168 if (item == NULL)
2169 return NULL;
2170
2171 if (lz->func == Py_None) {
2172 ok = PyObject_IsTrue(item);
2173 } else {
2174 PyObject *good;
2175 good = PyObject_CallFunctionObjArgs(lz->func,
2176 item, NULL);
2177 if (good == NULL) {
2178 Py_DECREF(item);
2179 return NULL;
2180 }
2181 ok = PyObject_IsTrue(good);
2182 Py_DECREF(good);
2183 }
2184 if (!ok)
2185 return item;
2186 Py_DECREF(item);
2187 }
2188}
2189
Raymond Hettinger60eca932003-02-09 06:40:58 +00002190PyDoc_STRVAR(ifilterfalse_doc,
2191"ifilterfalse(function or None, sequence) --> ifilterfalse object\n\
2192\n\
2193Return those items of sequence for which function(item) is false.\n\
2194If function is None, return the items that are false.");
2195
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002196static PyTypeObject ifilterfalse_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002197 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002198 "itertools.ifilterfalse", /* tp_name */
2199 sizeof(ifilterfalseobject), /* tp_basicsize */
2200 0, /* tp_itemsize */
2201 /* methods */
2202 (destructor)ifilterfalse_dealloc, /* tp_dealloc */
2203 0, /* tp_print */
2204 0, /* tp_getattr */
2205 0, /* tp_setattr */
2206 0, /* tp_compare */
2207 0, /* tp_repr */
2208 0, /* tp_as_number */
2209 0, /* tp_as_sequence */
2210 0, /* tp_as_mapping */
2211 0, /* tp_hash */
2212 0, /* tp_call */
2213 0, /* tp_str */
2214 PyObject_GenericGetAttr, /* tp_getattro */
2215 0, /* tp_setattro */
2216 0, /* tp_as_buffer */
2217 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2218 Py_TPFLAGS_BASETYPE, /* tp_flags */
2219 ifilterfalse_doc, /* tp_doc */
2220 (traverseproc)ifilterfalse_traverse, /* tp_traverse */
2221 0, /* tp_clear */
2222 0, /* tp_richcompare */
2223 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002224 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002225 (iternextfunc)ifilterfalse_next, /* tp_iternext */
2226 0, /* tp_methods */
2227 0, /* tp_members */
2228 0, /* tp_getset */
2229 0, /* tp_base */
2230 0, /* tp_dict */
2231 0, /* tp_descr_get */
2232 0, /* tp_descr_set */
2233 0, /* tp_dictoffset */
2234 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002235 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002236 ifilterfalse_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002237 PyObject_GC_Del, /* tp_free */
2238};
2239
2240
2241/* count object ************************************************************/
2242
2243typedef struct {
2244 PyObject_HEAD
Jack Diederich6c433a92006-05-26 11:15:17 +00002245 Py_ssize_t cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002246 PyObject *long_cnt; /* Arbitrarily large count when cnt >= PY_SSIZE_T_MAX */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002247} countobject;
2248
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002249static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002250
2251static PyObject *
2252count_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2253{
2254 countobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002255 Py_ssize_t cnt = 0;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002256 PyObject *cnt_arg = NULL;
2257 PyObject *long_cnt = NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002258
Georg Brandlb84c1372007-01-21 10:28:43 +00002259 if (type == &count_type && !_PyArg_NoKeywords("count()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002260 return NULL;
2261
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002262 if (!PyArg_UnpackTuple(args, "count", 0, 1, &cnt_arg))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002263 return NULL;
2264
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002265 if (cnt_arg != NULL) {
2266 cnt = PyInt_AsSsize_t(cnt_arg);
2267 if (cnt == -1 && PyErr_Occurred()) {
2268 PyErr_Clear();
2269 if (!PyLong_Check(cnt_arg)) {
2270 PyErr_SetString(PyExc_TypeError, "an integer is required");
2271 return NULL;
2272 }
2273 long_cnt = cnt_arg;
2274 Py_INCREF(long_cnt);
2275 cnt = PY_SSIZE_T_MAX;
2276 }
2277 }
2278
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002279 /* create countobject structure */
2280 lz = (countobject *)PyObject_New(countobject, &count_type);
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002281 if (lz == NULL) {
2282 Py_XDECREF(long_cnt);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002283 return NULL;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002284 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002285 lz->cnt = cnt;
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002286 lz->long_cnt = long_cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002287
2288 return (PyObject *)lz;
2289}
2290
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002291static void
2292count_dealloc(countobject *lz)
2293{
2294 Py_XDECREF(lz->long_cnt);
2295 PyObject_Del(lz);
2296}
2297
2298static PyObject *
2299count_nextlong(countobject *lz)
2300{
2301 static PyObject *one = NULL;
2302 PyObject *cnt;
2303 PyObject *stepped_up;
2304
2305 if (lz->long_cnt == NULL) {
2306 lz->long_cnt = PyInt_FromSsize_t(PY_SSIZE_T_MAX);
2307 if (lz->long_cnt == NULL)
2308 return NULL;
2309 }
2310 if (one == NULL) {
2311 one = PyInt_FromLong(1);
2312 if (one == NULL)
2313 return NULL;
2314 }
2315 cnt = lz->long_cnt;
2316 assert(cnt != NULL);
2317 stepped_up = PyNumber_Add(cnt, one);
2318 if (stepped_up == NULL)
2319 return NULL;
2320 lz->long_cnt = stepped_up;
2321 return cnt;
2322}
2323
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002324static PyObject *
2325count_next(countobject *lz)
2326{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002327 if (lz->cnt == PY_SSIZE_T_MAX)
2328 return count_nextlong(lz);
Jack Diederich36234e82006-09-21 17:50:26 +00002329 return PyInt_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002330}
2331
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002332static PyObject *
2333count_repr(countobject *lz)
2334{
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002335 PyObject *cnt_repr;
2336 PyObject *result;
2337
2338 if (lz->cnt != PY_SSIZE_T_MAX)
2339 return PyString_FromFormat("count(%zd)", lz->cnt);
2340
2341 cnt_repr = PyObject_Repr(lz->long_cnt);
2342 if (cnt_repr == NULL)
2343 return NULL;
2344 result = PyString_FromFormat("count(%s)", PyString_AS_STRING(cnt_repr));
2345 Py_DECREF(cnt_repr);
2346 return result;
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002347}
2348
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002349PyDoc_STRVAR(count_doc,
2350"count([firstval]) --> count object\n\
2351\n\
2352Return a count object whose .next() method returns consecutive\n\
2353integers starting from zero or, if specified, from firstval.");
2354
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002355static PyTypeObject count_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002356 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002357 "itertools.count", /* tp_name */
2358 sizeof(countobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002359 0, /* tp_itemsize */
2360 /* methods */
Raymond Hettinger50e90e22007-10-04 00:20:27 +00002361 (destructor)count_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002362 0, /* tp_print */
2363 0, /* tp_getattr */
2364 0, /* tp_setattr */
2365 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002366 (reprfunc)count_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002367 0, /* tp_as_number */
2368 0, /* tp_as_sequence */
2369 0, /* tp_as_mapping */
2370 0, /* tp_hash */
2371 0, /* tp_call */
2372 0, /* tp_str */
2373 PyObject_GenericGetAttr, /* tp_getattro */
2374 0, /* tp_setattro */
2375 0, /* tp_as_buffer */
2376 Py_TPFLAGS_DEFAULT, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002377 count_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002378 0, /* tp_traverse */
2379 0, /* tp_clear */
2380 0, /* tp_richcompare */
2381 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002382 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002383 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002384 0, /* tp_methods */
2385 0, /* tp_members */
2386 0, /* tp_getset */
2387 0, /* tp_base */
2388 0, /* tp_dict */
2389 0, /* tp_descr_get */
2390 0, /* tp_descr_set */
2391 0, /* tp_dictoffset */
2392 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002393 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002394 count_new, /* tp_new */
2395};
2396
2397
2398/* izip object ************************************************************/
2399
2400#include "Python.h"
2401
2402typedef struct {
2403 PyObject_HEAD
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002404 Py_ssize_t tuplesize;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002405 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002406 PyObject *result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002407} izipobject;
2408
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002409static PyTypeObject izip_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002410
2411static PyObject *
2412izip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2413{
2414 izipobject *lz;
Jack Diederich6c433a92006-05-26 11:15:17 +00002415 Py_ssize_t i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002416 PyObject *ittuple; /* tuple of iterators */
Raymond Hettinger2012f172003-02-07 05:32:58 +00002417 PyObject *result;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002418 Py_ssize_t tuplesize = PySequence_Length(args);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002419
Georg Brandlb84c1372007-01-21 10:28:43 +00002420 if (type == &izip_type && !_PyArg_NoKeywords("izip()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002421 return NULL;
2422
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002423 /* args must be a tuple */
2424 assert(PyTuple_Check(args));
2425
2426 /* obtain iterators */
2427 ittuple = PyTuple_New(tuplesize);
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002428 if (ittuple == NULL)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002429 return NULL;
2430 for (i=0; i < tuplesize; ++i) {
2431 PyObject *item = PyTuple_GET_ITEM(args, i);
2432 PyObject *it = PyObject_GetIter(item);
2433 if (it == NULL) {
2434 if (PyErr_ExceptionMatches(PyExc_TypeError))
2435 PyErr_Format(PyExc_TypeError,
Neal Norwitz2f3136b2006-05-27 05:18:57 +00002436 "izip argument #%zd must support iteration",
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002437 i+1);
2438 Py_DECREF(ittuple);
2439 return NULL;
2440 }
2441 PyTuple_SET_ITEM(ittuple, i, it);
2442 }
2443
Raymond Hettinger2012f172003-02-07 05:32:58 +00002444 /* create a result holder */
2445 result = PyTuple_New(tuplesize);
2446 if (result == NULL) {
2447 Py_DECREF(ittuple);
2448 return NULL;
2449 }
2450 for (i=0 ; i < tuplesize ; i++) {
2451 Py_INCREF(Py_None);
2452 PyTuple_SET_ITEM(result, i, Py_None);
2453 }
2454
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002455 /* create izipobject structure */
2456 lz = (izipobject *)type->tp_alloc(type, 0);
2457 if (lz == NULL) {
2458 Py_DECREF(ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002459 Py_DECREF(result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002460 return NULL;
2461 }
2462 lz->ittuple = ittuple;
2463 lz->tuplesize = tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002464 lz->result = result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002465
2466 return (PyObject *)lz;
2467}
2468
2469static void
2470izip_dealloc(izipobject *lz)
2471{
2472 PyObject_GC_UnTrack(lz);
2473 Py_XDECREF(lz->ittuple);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002474 Py_XDECREF(lz->result);
Christian Heimese93237d2007-12-19 02:37:44 +00002475 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002476}
2477
2478static int
2479izip_traverse(izipobject *lz, visitproc visit, void *arg)
2480{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002481 Py_VISIT(lz->ittuple);
2482 Py_VISIT(lz->result);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002483 return 0;
2484}
2485
2486static PyObject *
2487izip_next(izipobject *lz)
2488{
Jack Diederich6c433a92006-05-26 11:15:17 +00002489 Py_ssize_t i;
Martin v. Löwisad0a4622006-02-16 14:30:23 +00002490 Py_ssize_t tuplesize = lz->tuplesize;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002491 PyObject *result = lz->result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002492 PyObject *it;
2493 PyObject *item;
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002494 PyObject *olditem;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002495
Raymond Hettingerb5a42082003-08-08 05:10:41 +00002496 if (tuplesize == 0)
2497 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00002498 if (Py_REFCNT(result) == 1) {
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002499 Py_INCREF(result);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002500 for (i=0 ; i < tuplesize ; i++) {
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002501 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002502 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002503 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002504 if (item == NULL) {
2505 Py_DECREF(result);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002506 return NULL;
Raymond Hettingera56f6b62003-08-29 23:09:58 +00002507 }
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002508 olditem = PyTuple_GET_ITEM(result, i);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002509 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger4f01f892003-08-30 00:10:06 +00002510 Py_DECREF(olditem);
Raymond Hettinger2012f172003-02-07 05:32:58 +00002511 }
Raymond Hettinger2012f172003-02-07 05:32:58 +00002512 } else {
Raymond Hettinger2012f172003-02-07 05:32:58 +00002513 result = PyTuple_New(tuplesize);
2514 if (result == NULL)
2515 return NULL;
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002516 for (i=0 ; i < tuplesize ; i++) {
2517 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettingerd1a283b2003-03-01 01:48:24 +00002518 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002519 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettingerf0c00242003-02-07 07:26:25 +00002520 if (item == NULL) {
2521 Py_DECREF(result);
2522 return NULL;
2523 }
2524 PyTuple_SET_ITEM(result, i, item);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002525 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002526 }
2527 return result;
2528}
2529
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002530PyDoc_STRVAR(izip_doc,
2531"izip(iter1 [,iter2 [...]]) --> izip object\n\
2532\n\
2533Return a izip object whose .next() method returns a tuple where\n\
2534the i-th element comes from the i-th iterable argument. The .next()\n\
2535method continues until the shortest iterable in the argument sequence\n\
2536is exhausted and then it raises StopIteration. Works like the zip()\n\
2537function but consumes less memory by returning an iterator instead of\n\
2538a list.");
2539
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002540static PyTypeObject izip_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002541 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger60eca932003-02-09 06:40:58 +00002542 "itertools.izip", /* tp_name */
2543 sizeof(izipobject), /* tp_basicsize */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002544 0, /* tp_itemsize */
2545 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002546 (destructor)izip_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002547 0, /* tp_print */
2548 0, /* tp_getattr */
2549 0, /* tp_setattr */
2550 0, /* tp_compare */
2551 0, /* tp_repr */
2552 0, /* tp_as_number */
2553 0, /* tp_as_sequence */
2554 0, /* tp_as_mapping */
2555 0, /* tp_hash */
2556 0, /* tp_call */
2557 0, /* tp_str */
2558 PyObject_GenericGetAttr, /* tp_getattro */
2559 0, /* tp_setattro */
2560 0, /* tp_as_buffer */
2561 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2562 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002563 izip_doc, /* tp_doc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002564 (traverseproc)izip_traverse, /* tp_traverse */
2565 0, /* tp_clear */
2566 0, /* tp_richcompare */
2567 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002568 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002569 (iternextfunc)izip_next, /* tp_iternext */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002570 0, /* tp_methods */
2571 0, /* tp_members */
2572 0, /* tp_getset */
2573 0, /* tp_base */
2574 0, /* tp_dict */
2575 0, /* tp_descr_get */
2576 0, /* tp_descr_set */
2577 0, /* tp_dictoffset */
2578 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002579 0, /* tp_alloc */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002580 izip_new, /* tp_new */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002581 PyObject_GC_Del, /* tp_free */
2582};
2583
2584
2585/* repeat object ************************************************************/
2586
2587typedef struct {
2588 PyObject_HEAD
2589 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002590 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002591} repeatobject;
2592
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002593static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002594
2595static PyObject *
2596repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2597{
2598 repeatobject *ro;
2599 PyObject *element;
Jack Diederich6c433a92006-05-26 11:15:17 +00002600 Py_ssize_t cnt = -1;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002601
Georg Brandlb84c1372007-01-21 10:28:43 +00002602 if (type == &repeat_type && !_PyArg_NoKeywords("repeat()", kwds))
Georg Brandl02c42872005-08-26 06:42:30 +00002603 return NULL;
2604
Jack Diederich6c433a92006-05-26 11:15:17 +00002605 if (!PyArg_ParseTuple(args, "O|n:repeat", &element, &cnt))
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002606 return NULL;
2607
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00002608 if (PyTuple_Size(args) == 2 && cnt < 0)
2609 cnt = 0;
2610
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002611 ro = (repeatobject *)type->tp_alloc(type, 0);
2612 if (ro == NULL)
2613 return NULL;
2614 Py_INCREF(element);
2615 ro->element = element;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002616 ro->cnt = cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002617 return (PyObject *)ro;
2618}
2619
2620static void
2621repeat_dealloc(repeatobject *ro)
2622{
2623 PyObject_GC_UnTrack(ro);
2624 Py_XDECREF(ro->element);
Christian Heimese93237d2007-12-19 02:37:44 +00002625 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002626}
2627
2628static int
2629repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
2630{
Raymond Hettinger58ed69b2004-07-15 05:32:47 +00002631 Py_VISIT(ro->element);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002632 return 0;
2633}
2634
2635static PyObject *
2636repeat_next(repeatobject *ro)
2637{
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002638 if (ro->cnt == 0)
2639 return NULL;
2640 if (ro->cnt > 0)
2641 ro->cnt--;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002642 Py_INCREF(ro->element);
2643 return ro->element;
2644}
2645
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002646static PyObject *
2647repeat_repr(repeatobject *ro)
2648{
2649 PyObject *result, *objrepr;
2650
2651 objrepr = PyObject_Repr(ro->element);
2652 if (objrepr == NULL)
2653 return NULL;
2654
2655 if (ro->cnt == -1)
2656 result = PyString_FromFormat("repeat(%s)",
2657 PyString_AS_STRING(objrepr));
2658 else
Jack Diederich6c433a92006-05-26 11:15:17 +00002659 result = PyString_FromFormat("repeat(%s, %zd)",
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002660 PyString_AS_STRING(objrepr), ro->cnt);
2661 Py_DECREF(objrepr);
2662 return result;
2663}
2664
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002665static PyObject *
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002666repeat_len(repeatobject *ro)
2667{
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002668 if (ro->cnt == -1) {
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002669 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002670 return NULL;
2671 }
Jack Diederich6c433a92006-05-26 11:15:17 +00002672 return PyInt_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002673}
2674
Armin Rigof5b3e362006-02-11 21:32:43 +00002675PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002676
2677static PyMethodDef repeat_methods[] = {
Armin Rigof5b3e362006-02-11 21:32:43 +00002678 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002679 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00002680};
2681
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002682PyDoc_STRVAR(repeat_doc,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002683"repeat(element [,times]) -> create an iterator which returns the element\n\
2684for the specified number of times. If not specified, returns the element\n\
2685endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002686
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002687static PyTypeObject repeat_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002688 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002689 "itertools.repeat", /* tp_name */
2690 sizeof(repeatobject), /* tp_basicsize */
2691 0, /* tp_itemsize */
2692 /* methods */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002693 (destructor)repeat_dealloc, /* tp_dealloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002694 0, /* tp_print */
2695 0, /* tp_getattr */
2696 0, /* tp_setattr */
2697 0, /* tp_compare */
Raymond Hettinger7dacda22004-04-08 21:54:00 +00002698 (reprfunc)repeat_repr, /* tp_repr */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002699 0, /* tp_as_number */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002700 0, /* tp_as_sequence */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002701 0, /* tp_as_mapping */
2702 0, /* tp_hash */
2703 0, /* tp_call */
2704 0, /* tp_str */
2705 PyObject_GenericGetAttr, /* tp_getattro */
2706 0, /* tp_setattro */
2707 0, /* tp_as_buffer */
2708 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2709 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002710 repeat_doc, /* tp_doc */
2711 (traverseproc)repeat_traverse, /* tp_traverse */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002712 0, /* tp_clear */
2713 0, /* tp_richcompare */
2714 0, /* tp_weaklistoffset */
Raymond Hettinger1da1dbf2003-03-17 19:46:11 +00002715 PyObject_SelfIter, /* tp_iter */
Raymond Hettinger60eca932003-02-09 06:40:58 +00002716 (iternextfunc)repeat_next, /* tp_iternext */
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00002717 repeat_methods, /* tp_methods */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002718 0, /* tp_members */
2719 0, /* tp_getset */
2720 0, /* tp_base */
2721 0, /* tp_dict */
2722 0, /* tp_descr_get */
2723 0, /* tp_descr_set */
2724 0, /* tp_dictoffset */
2725 0, /* tp_init */
Raymond Hettingerbfef18c2003-05-23 03:55:42 +00002726 0, /* tp_alloc */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002727 repeat_new, /* tp_new */
2728 PyObject_GC_Del, /* tp_free */
2729};
2730
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002731/* iziplongest object ************************************************************/
2732
2733#include "Python.h"
2734
2735typedef struct {
2736 PyObject_HEAD
2737 Py_ssize_t tuplesize;
2738 Py_ssize_t numactive;
2739 PyObject *ittuple; /* tuple of iterators */
2740 PyObject *result;
2741 PyObject *fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002742} iziplongestobject;
2743
2744static PyTypeObject iziplongest_type;
2745
2746static PyObject *
2747izip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2748{
2749 iziplongestobject *lz;
2750 Py_ssize_t i;
2751 PyObject *ittuple; /* tuple of iterators */
2752 PyObject *result;
2753 PyObject *fillvalue = Py_None;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002754 Py_ssize_t tuplesize = PySequence_Length(args);
2755
2756 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_Size(kwds) > 0) {
2757 fillvalue = PyDict_GetItemString(kwds, "fillvalue");
2758 if (fillvalue == NULL || PyDict_Size(kwds) > 1) {
2759 PyErr_SetString(PyExc_TypeError,
2760 "izip_longest() got an unexpected keyword argument");
2761 return NULL;
2762 }
2763 }
2764
2765 /* args must be a tuple */
2766 assert(PyTuple_Check(args));
2767
2768 /* obtain iterators */
2769 ittuple = PyTuple_New(tuplesize);
2770 if (ittuple == NULL)
2771 return NULL;
2772 for (i=0; i < tuplesize; ++i) {
2773 PyObject *item = PyTuple_GET_ITEM(args, i);
2774 PyObject *it = PyObject_GetIter(item);
2775 if (it == NULL) {
2776 if (PyErr_ExceptionMatches(PyExc_TypeError))
2777 PyErr_Format(PyExc_TypeError,
2778 "izip_longest argument #%zd must support iteration",
2779 i+1);
2780 Py_DECREF(ittuple);
2781 return NULL;
2782 }
2783 PyTuple_SET_ITEM(ittuple, i, it);
2784 }
2785
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002786 /* create a result holder */
2787 result = PyTuple_New(tuplesize);
2788 if (result == NULL) {
2789 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002790 return NULL;
2791 }
2792 for (i=0 ; i < tuplesize ; i++) {
2793 Py_INCREF(Py_None);
2794 PyTuple_SET_ITEM(result, i, Py_None);
2795 }
2796
2797 /* create iziplongestobject structure */
2798 lz = (iziplongestobject *)type->tp_alloc(type, 0);
2799 if (lz == NULL) {
2800 Py_DECREF(ittuple);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002801 Py_DECREF(result);
2802 return NULL;
2803 }
2804 lz->ittuple = ittuple;
2805 lz->tuplesize = tuplesize;
2806 lz->numactive = tuplesize;
2807 lz->result = result;
2808 Py_INCREF(fillvalue);
2809 lz->fillvalue = fillvalue;
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002810 return (PyObject *)lz;
2811}
2812
2813static void
2814izip_longest_dealloc(iziplongestobject *lz)
2815{
2816 PyObject_GC_UnTrack(lz);
2817 Py_XDECREF(lz->ittuple);
2818 Py_XDECREF(lz->result);
2819 Py_XDECREF(lz->fillvalue);
Christian Heimese93237d2007-12-19 02:37:44 +00002820 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002821}
2822
2823static int
2824izip_longest_traverse(iziplongestobject *lz, visitproc visit, void *arg)
2825{
2826 Py_VISIT(lz->ittuple);
2827 Py_VISIT(lz->result);
2828 Py_VISIT(lz->fillvalue);
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002829 return 0;
2830}
2831
2832static PyObject *
2833izip_longest_next(iziplongestobject *lz)
2834{
2835 Py_ssize_t i;
2836 Py_ssize_t tuplesize = lz->tuplesize;
2837 PyObject *result = lz->result;
2838 PyObject *it;
2839 PyObject *item;
2840 PyObject *olditem;
2841
2842 if (tuplesize == 0)
2843 return NULL;
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002844 if (lz->numactive == 0)
2845 return NULL;
Christian Heimese93237d2007-12-19 02:37:44 +00002846 if (Py_REFCNT(result) == 1) {
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002847 Py_INCREF(result);
2848 for (i=0 ; i < tuplesize ; i++) {
2849 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002850 if (it == NULL) {
2851 Py_INCREF(lz->fillvalue);
2852 item = lz->fillvalue;
2853 } else {
2854 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002855 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002856 if (item == NULL) {
2857 lz->numactive -= 1;
2858 if (lz->numactive == 0) {
2859 Py_DECREF(result);
2860 return NULL;
2861 } else {
2862 Py_INCREF(lz->fillvalue);
2863 item = lz->fillvalue;
2864 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
2865 Py_DECREF(it);
2866 }
2867 }
2868 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002869 olditem = PyTuple_GET_ITEM(result, i);
2870 PyTuple_SET_ITEM(result, i, item);
2871 Py_DECREF(olditem);
2872 }
2873 } else {
2874 result = PyTuple_New(tuplesize);
2875 if (result == NULL)
2876 return NULL;
2877 for (i=0 ; i < tuplesize ; i++) {
2878 it = PyTuple_GET_ITEM(lz->ittuple, i);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002879 if (it == NULL) {
2880 Py_INCREF(lz->fillvalue);
2881 item = lz->fillvalue;
2882 } else {
2883 assert(PyIter_Check(it));
Christian Heimese93237d2007-12-19 02:37:44 +00002884 item = (*Py_TYPE(it)->tp_iternext)(it);
Raymond Hettinger1b6ca542007-02-21 17:22:05 +00002885 if (item == NULL) {
2886 lz->numactive -= 1;
2887 if (lz->numactive == 0) {
2888 Py_DECREF(result);
2889 return NULL;
2890 } else {
2891 Py_INCREF(lz->fillvalue);
2892 item = lz->fillvalue;
2893 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
2894 Py_DECREF(it);
2895 }
2896 }
2897 }
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002898 PyTuple_SET_ITEM(result, i, item);
2899 }
2900 }
2901 return result;
2902}
2903
2904PyDoc_STRVAR(izip_longest_doc,
2905"izip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> izip_longest object\n\
2906\n\
2907Return an izip_longest object whose .next() method returns a tuple where\n\
2908the i-th element comes from the i-th iterable argument. The .next()\n\
2909method continues until the longest iterable in the argument sequence\n\
2910is exhausted and then it raises StopIteration. When the shorter iterables\n\
2911are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
2912defaults to None or can be specified by a keyword argument.\n\
2913");
2914
2915static PyTypeObject iziplongest_type = {
Martin v. Löwis68192102007-07-21 06:55:02 +00002916 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002917 "itertools.izip_longest", /* tp_name */
2918 sizeof(iziplongestobject), /* tp_basicsize */
2919 0, /* tp_itemsize */
2920 /* methods */
2921 (destructor)izip_longest_dealloc, /* tp_dealloc */
2922 0, /* tp_print */
2923 0, /* tp_getattr */
2924 0, /* tp_setattr */
2925 0, /* tp_compare */
2926 0, /* tp_repr */
2927 0, /* tp_as_number */
2928 0, /* tp_as_sequence */
2929 0, /* tp_as_mapping */
2930 0, /* tp_hash */
2931 0, /* tp_call */
2932 0, /* tp_str */
2933 PyObject_GenericGetAttr, /* tp_getattro */
2934 0, /* tp_setattro */
2935 0, /* tp_as_buffer */
2936 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2937 Py_TPFLAGS_BASETYPE, /* tp_flags */
2938 izip_longest_doc, /* tp_doc */
2939 (traverseproc)izip_longest_traverse, /* tp_traverse */
2940 0, /* tp_clear */
2941 0, /* tp_richcompare */
2942 0, /* tp_weaklistoffset */
2943 PyObject_SelfIter, /* tp_iter */
2944 (iternextfunc)izip_longest_next, /* tp_iternext */
2945 0, /* tp_methods */
2946 0, /* tp_members */
2947 0, /* tp_getset */
2948 0, /* tp_base */
2949 0, /* tp_dict */
2950 0, /* tp_descr_get */
2951 0, /* tp_descr_set */
2952 0, /* tp_dictoffset */
2953 0, /* tp_init */
2954 0, /* tp_alloc */
2955 izip_longest_new, /* tp_new */
2956 PyObject_GC_Del, /* tp_free */
2957};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002958
2959/* module level code ********************************************************/
2960
2961PyDoc_STRVAR(module_doc,
2962"Functional tools for creating and using iterators.\n\
2963\n\
2964Infinite iterators:\n\
2965count([n]) --> n, n+1, n+2, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002966cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00002967repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002968\n\
2969Iterators terminating on the shortest input sequence:\n\
2970izip(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettingerd36862c2007-02-21 05:20:38 +00002971izip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... \n\
Raymond Hettinger60eca932003-02-09 06:40:58 +00002972ifilter(pred, seq) --> elements of seq where pred(elem) is True\n\
2973ifilterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002974islice(seq, [start,] stop [, step]) --> elements from\n\
2975 seq[start:stop:step]\n\
2976imap(fun, p, q, ...) --> fun(p0, q0), fun(p1, q1), ...\n\
2977starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00002978tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002979chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002980takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
2981dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00002982groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002983");
2984
2985
Raymond Hettingerad983e72003-11-12 14:32:26 +00002986static PyMethodDef module_methods[] = {
2987 {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc},
2988 {NULL, NULL} /* sentinel */
2989};
2990
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002991PyMODINIT_FUNC
2992inititertools(void)
2993{
Raymond Hettinger60eca932003-02-09 06:40:58 +00002994 int i;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002995 PyObject *m;
Raymond Hettinger60eca932003-02-09 06:40:58 +00002996 char *name;
2997 PyTypeObject *typelist[] = {
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002998 &cycle_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00002999 &dropwhile_type,
3000 &takewhile_type,
3001 &islice_type,
3002 &starmap_type,
3003 &imap_type,
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003004 &chain_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003005 &ifilter_type,
3006 &ifilterfalse_type,
3007 &count_type,
3008 &izip_type,
Raymond Hettinger50986cc2008-02-22 03:16:42 +00003009 &iziplongest_type,
3010 &product_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003011 &repeat_type,
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003012 &groupby_type,
Raymond Hettinger60eca932003-02-09 06:40:58 +00003013 NULL
3014 };
3015
Christian Heimese93237d2007-12-19 02:37:44 +00003016 Py_TYPE(&teedataobject_type) = &PyType_Type;
Raymond Hettingerad983e72003-11-12 14:32:26 +00003017 m = Py_InitModule3("itertools", module_methods, module_doc);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00003018 if (m == NULL)
3019 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003020
Raymond Hettinger60eca932003-02-09 06:40:58 +00003021 for (i=0 ; typelist[i] != NULL ; i++) {
3022 if (PyType_Ready(typelist[i]) < 0)
3023 return;
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003024 name = strchr(typelist[i]->tp_name, '.');
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00003025 assert (name != NULL);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003026 Py_INCREF(typelist[i]);
Raymond Hettingerd449eab2003-05-22 16:32:58 +00003027 PyModule_AddObject(m, name+1, (PyObject *)typelist[i]);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003028 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00003029
3030 if (PyType_Ready(&teedataobject_type) < 0)
3031 return;
3032 if (PyType_Ready(&tee_type) < 0)
3033 return;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +00003034 if (PyType_Ready(&_grouper_type) < 0)
3035 return;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003036}